use std::fmt;
use std::io::Result;
#[cfg(feature = "wal")]
use std::sync::Arc;
#[cfg(feature = "wal")]
use std::sync::Weak;
#[cfg(feature = "wal")]
use std::pin::Pin;
#[cfg(not(feature = "wal"))]
use bytes::Bytes;
use crate::SegmentId;
#[cfg(feature = "wal")]
use crate::Cno;
use crate::ondisk::{SegmentHeader, SegmentBlockEntryRaw};
use crate::{BlockIndex, BlockPtr};
use crate::config::HyperFileMetaConfig;
use crate::inode::Inode;
use crate::ondisk::InodeRaw;
use crate::buffer::DataBlock;
use crate::segment_body::SegmentBody;
pub struct Segment;
impl Segment {
pub fn segid_to_staging_file_id(segid: SegmentId) -> String {
match segid.part_id() {
None => format!("{:0>10}", segid.seq_id()),
Some(p) => format!("{:0>10}.{}", segid.seq_id(), p),
}
}
pub const SUMMARY_PART: u16 = 0;
pub const FIRST_DATA_PART: u16 = 1;
pub const MAX_PARTS: usize = crate::meta_format::BLOCK_PTR_PARTED_MAX_PARTS;
pub fn object_of(segid: SegmentId, config: &HyperFileMetaConfig) -> SegmentId {
if config.block_ptr_format.is_parted() {
segid.at_part(Self::SUMMARY_PART)
} else {
segid.checkpoint()
}
}
}
pub trait SegmentReadWrite {
fn append(&self, segid: SegmentId, buf: &[u8]) -> Result<()>;
fn done(&self, segid: SegmentId, buf: &[u8], len: usize) -> impl Future<Output = Result<()>> + Send;
fn done_pieces(&self, segid: SegmentId, body: SegmentBody) -> impl Future<Output = Result<()>> + Send;
fn remove(&self, segid: SegmentId) -> impl Future<Output = Result<()>>;
fn open(&self, segid: SegmentId) -> impl Future<Output = Result<SegmentSum>>;
fn list(&self, segid: SegmentId) -> impl Future<Output = Result<Vec<SegmentId>>>;
fn build_block_map(&self, segid: SegmentId) -> impl Future<Output = Result<Vec<(BlockIndex, BlockPtr)>>>;
}
const SEGMENT_SUMMARY_HEADER_SIZE: usize = std::mem::size_of::<SegmentHeader>();
const SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE: usize = std::mem::size_of::<SegmentBlockEntryRaw>();
#[derive(Debug)]
pub struct SegmentBlockDesc {
pub(crate) blkidx: BlockIndex,
pub(crate) blkptr: BlockPtr,
}
impl fmt::Display for SegmentBlockDesc {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "SegmentBlockDesc {{ blkidx: {}, blkptr: {} }}", self.blkidx, self.blkptr)
}
}
pub struct SegmentSum {
pub hdr: SegmentHeader,
pub blocks: Vec<SegmentBlockDesc>,
}
impl fmt::Display for SegmentSum {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
writeln!(f, "==== dump Segment Summary ====")?;
writeln!(f, " ino: {}, cno: {}, built on inode {:#018x}", self.hdr.s_ino, self.hdr.s_cno, self.hdr.s_base)?;
writeln!(f, " magic {:#x}, checksum: {:#x}", self.hdr.s_magic, self.hdr.s_chksum)?;
writeln!(f, " bytes: {}, flags: {:#x}", self.hdr.s_bytes, self.hdr.s_flags)?;
writeln!(f, " meta block shift: {}, data block shift: {}", self.hdr.s_meta_blk_shift, self.hdr.s_data_blk_shift)?;
writeln!(f, " meta block count: {}, data block count: {}", self.hdr.s_nmetablk, self.hdr.s_ndatablk)?;
for blkdesc in self.blocks.iter() {
writeln!(f, " {}", blkdesc)?;
}
writeln!(f, " {}", Inode::from_raw(&self.hdr.s_inode, None))
}
}
impl SegmentSum {
pub fn from_slice(buf: &[u8]) -> Self {
let hdrsz = SEGMENT_SUMMARY_HEADER_SIZE;
let bufsz = buf.len();
if bufsz < hdrsz {
panic!("failed to create segment sum from buf, buf len {} < hdr size {}, it's too small", bufsz, hdrsz);
}
let hdr = SegmentHeader::read_from(&buf[..hdrsz]);
let ndatablk = hdr.s_ndatablk as usize;
let need = hdrsz + ndatablk * SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
if bufsz < need {
panic!("failed to create segment sum from buf, buf len {} < hdr size {} + num of blk idx {}, it's too small", bufsz, hdrsz, ndatablk);
}
let mut blocks: Vec<SegmentBlockDesc> = Vec::with_capacity(ndatablk);
let mut off = hdrsz;
for _ in 0..ndatablk {
let raw = SegmentBlockEntryRaw::read_from(&buf[off..off + SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE]);
blocks.push(SegmentBlockDesc {
blkidx: raw.e_blkidx,
blkptr: raw.e_blkptr,
});
off += SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
}
Self { hdr, blocks }
}
pub fn write_to(&self, buf: &mut [u8]) {
let hdrsz = SEGMENT_SUMMARY_HEADER_SIZE;
let bufsz = buf.len();
if bufsz < hdrsz {
panic!("failed to write segment sum to buf, buf len {} < hdr size {}, it's too small", bufsz, hdrsz);
}
let ndatablk = self.hdr.s_ndatablk as usize;
let need = hdrsz + ndatablk * SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
if bufsz < need {
panic!("failed to write segment sum to buf, buf len {} < hdr size {} + num of blk idx {}, it's too small", bufsz, hdrsz, ndatablk);
}
self.hdr.write_to(&mut buf[..hdrsz]);
let mut off = hdrsz;
for entry in self.blocks.iter() {
let raw = SegmentBlockEntryRaw {
e_blkidx: entry.blkidx,
e_blkptr: entry.blkptr,
};
raw.write_to(&mut buf[off..off + SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE]);
off += SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
}
}
pub fn update(&mut self, chksum: u32, inode: &InodeRaw) -> usize {
let data_block_count = self.blocks.len();
let ss_bytes = SEGMENT_SUMMARY_HEADER_SIZE + data_block_count * SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
assert!(self.hdr.s_ndatablk == data_block_count as u32);
self.hdr.s_inode = inode.to_owned();
self.hdr.s_bytes = ss_bytes as u32;
self.hdr.s_chksum = chksum;
ss_bytes
}
pub fn calc_staging_off(&self, data_block_index: usize) -> usize {
let mut offset = 0;
offset += self.hdr.aligned_ss_bytes();
offset += (self.hdr.s_nmetablk << self.hdr.s_meta_blk_shift) as usize;
offset += data_block_index << self.hdr.s_data_blk_shift;
offset
}
}
pub struct Writer<T> {
ctx: T,
#[cfg(not(feature = "wal"))]
header: Vec<u8>,
#[cfg(not(feature = "wal"))]
body: SegmentBody,
#[cfg(feature = "wal")]
data: Arc<Pin<Box<Vec<u8>>>>,
offset: usize,
segid: SegmentId,
ss: SegmentSum,
}
impl<T: SegmentReadWrite> Writer<T> {
pub fn new_at(ctx: T, buf_size: usize, segid: SegmentId, hyper_file_config: &HyperFileMetaConfig) -> Self {
let mut w = Self::new(ctx, buf_size, segid, hyper_file_config);
w.segid = segid;
w.ss.hdr.s_cno = segid.as_cno();
w
}
pub fn set_base(&mut self, fingerprint: u64) {
self.ss.hdr.s_base = fingerprint;
}
pub fn mark_partial(&mut self) {
self.ss.hdr.s_flags |= crate::ondisk::SEGMENT_FLAG_PARTIAL;
}
pub fn new(ctx: T, buf_size: usize, segid: SegmentId, hyper_file_config: &HyperFileMetaConfig) -> Self {
let _ = buf_size;
let mut hdr = SegmentHeader::new();
hdr.s_meta_blk_shift = hyper_file_config.meta_block_size.checked_ilog2().unwrap() as u8;
hdr.s_data_blk_shift = hyper_file_config.data_block_size.checked_ilog2().unwrap() as u8;
hdr.s_base = 0;
hdr.s_ino = 0;
hdr.s_cno = segid.as_cno();
Self {
ctx: ctx,
#[cfg(not(feature = "wal"))]
header: Vec::new(),
#[cfg(not(feature = "wal"))]
body: SegmentBody::new(),
#[cfg(feature = "wal")]
data: Arc::new(Box::pin(Vec::with_capacity(buf_size))),
offset: 0,
segid: Segment::object_of(segid, hyper_file_config),
ss: SegmentSum {
hdr: hdr,
blocks: Vec::new(),
},
}
}
#[inline]
pub fn calc_ss_aligned_bytes(ndatablk: usize) -> usize {
let ss_bytes = SEGMENT_SUMMARY_HEADER_SIZE + ndatablk * SEGMENT_SUMMARY_BLOCK_ENTRY_SIZE;
(ss_bytes + 4096 - 1) >> 12 << 12
}
pub fn realize_ss(&mut self, checksum: u32, inode: &InodeRaw) {
let ss_bytes = self.ss.update(checksum, inode);
#[cfg(not(feature = "wal"))]
let ss_aligned_bytes = {
self.header.resize(ss_bytes, 0);
self.ss.write_to(&mut self.header);
let ss_aligned_bytes = (ss_bytes + 4096 - 1) >> 12 << 12;
self.header.resize(ss_aligned_bytes, 0);
ss_aligned_bytes
};
#[cfg(feature = "wal")]
let ss_aligned_bytes = {
let Some(data) = Arc::get_mut(&mut self.data) else {
panic!("failed to get back inner data buffer during segment build");
};
data.resize(ss_bytes, 0);
self.ss.write_to(data);
let ss_aligned_bytes = (ss_bytes + 4096 - 1) >> 12 << 12;
data.resize(ss_aligned_bytes, 0);
ss_aligned_bytes
};
self.offset += ss_aligned_bytes;
}
pub fn append(&mut self, buf: &[u8]) -> Result<()> {
let len = buf.len();
#[cfg(not(feature = "wal"))]
{
self.body.push(Bytes::copy_from_slice(buf));
}
#[cfg(feature = "wal")]
let slice_start = {
let Some(data) = Arc::get_mut(&mut self.data) else {
panic!("failed to get back inner data buffer during segment build");
};
data.resize(self.offset + len, 0);
let slice_start = &mut data[self.offset..];
slice_start
};
#[cfg(feature = "wal")]
{
let (data, _) = slice_start.split_at_mut(len);
data.copy_from_slice(buf);
}
self.offset += len;
self.ctx.append(self.segid, buf)
}
pub fn append_data_block(&mut self, block: &DataBlock) -> Result<()> {
#[cfg(not(feature = "wal"))]
{
let bytes = block.bytes_view();
let len = bytes.len();
self.body.push(bytes);
self.offset += len;
self.ctx.append(self.segid, block.as_slice())
}
#[cfg(feature = "wal")]
{
self.append(block.as_slice())
}
}
#[cfg(all(feature = "concurrent-segment-build", feature = "wal"))]
pub fn spawn_append(&mut self, chunk: Vec<&DataBlock>) -> Result<tokio::task::JoinHandle<()>> {
let len = chunk.iter().map(|block| block.size()).sum();
let slice_start = {
let Some(data) = Arc::get_mut(&mut self.data) else {
panic!("failed to get back inner data buffer during segment build");
};
data.resize(self.offset + len, 0);
let slice_start = &mut data[self.offset..];
slice_start
};
let (data, _) = slice_start.split_at_mut(len);
let new_data = unsafe {
std::slice::from_raw_parts_mut(data.as_mut_ptr() as *mut u8, data.len())
};
let new_chunk = chunk.into_iter()
.map(|block| unsafe {
let buf = block.as_slice();
std::slice::from_raw_parts(buf.as_ptr() as *const u8, buf.len())
})
.collect::<Vec<&[u8]>>();
let join = tokio::task::spawn_blocking(move || {
for (i, buf) in new_chunk.into_iter().enumerate() {
let start = i * buf.len();
let end = start + buf.len();
let data = &mut new_data[start..end];
data.copy_from_slice(buf);
}
});
self.offset += len;
Ok(join)
}
pub async fn done(self) -> Result<()> {
#[cfg(not(feature = "wal"))]
{
let mut body = SegmentBody::with_capacity(self.body.pieces().len() + 1);
body.push(Bytes::from(self.header));
for piece in self.body.pieces() {
body.push(piece.clone());
}
assert_eq!(
body.len(),
self.offset,
"scatter body length must match the cumulative writer offset",
);
self.ctx.done_pieces(self.segid, body).await
}
#[cfg(feature = "wal")]
{
let _ = self.ctx.done(self.segid, &self.data, self.offset).await?;
unsafe { Arc::increment_strong_count(Arc::as_ptr(&self.data)) };
Ok(())
}
}
#[cfg(feature = "wal")]
pub fn get_weak_data(&self) -> (Cno, Weak<Pin<Box<Vec<u8>>>>) {
(self.segid.as_cno(), Arc::downgrade(&self.data))
}
}
impl<T> Writer<T> {
pub fn inc_metablk(&mut self) {
self.ss.hdr.s_nmetablk += 1;
}
pub fn inc_datablk(&mut self, blkidx: &BlockIndex, blkptr: &BlockPtr) {
self.ss.hdr.s_ndatablk += 1;
self.ss.blocks.push(SegmentBlockDesc { blkidx: *blkidx, blkptr: *blkptr });
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_segment_sum(ndatablk: usize) -> SegmentSum {
let mut hdr = SegmentHeader::new();
hdr.s_cno = 7;
hdr.s_ino = 1;
hdr.s_meta_blk_shift = 12;
hdr.s_data_blk_shift = 12;
hdr.s_nmetablk = 2;
hdr.s_ndatablk = ndatablk as u32;
let blocks: Vec<SegmentBlockDesc> = (0..ndatablk)
.map(|i| SegmentBlockDesc { blkidx: i as u64, blkptr: 0x4000_0001_0000_0000 + i as u64 })
.collect();
SegmentSum { hdr, blocks }
}
#[test]
fn segment_sum_write_read_round_trip() {
let ndatablk = 3;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0xDEAD, &inode);
let mut buf = vec![0u8; ss_bytes];
ss.write_to(&mut buf);
let restored = SegmentSum::from_slice(&buf);
assert_eq!(restored.hdr.s_magic, 0x48465353);
assert_eq!(restored.hdr.s_cno, 7);
assert_eq!(restored.hdr.s_nmetablk, 2);
assert_eq!(restored.hdr.s_ndatablk, ndatablk as u32);
assert_eq!(restored.hdr.s_chksum, 0xDEAD);
assert_eq!(restored.hdr.s_bytes, ss_bytes as u32);
assert_eq!(restored.blocks.len(), ndatablk);
for i in 0..ndatablk {
assert_eq!(restored.blocks[i].blkidx, i as u64);
assert_eq!(restored.blocks[i].blkptr, 0x4000_0001_0000_0000 + i as u64);
}
}
#[test]
fn segment_sum_zero_data_blocks() {
let mut ss = make_segment_sum(0);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0, &inode);
let mut buf = vec![0u8; ss_bytes];
ss.write_to(&mut buf);
let restored = SegmentSum::from_slice(&buf);
assert_eq!(restored.hdr.s_ndatablk, 0);
assert!(restored.blocks.is_empty());
}
#[test]
fn segment_sum_calc_staging_off() {
let ss = make_segment_sum(2);
let aligned = ss.hdr.aligned_ss_bytes();
assert_eq!(ss.calc_staging_off(0), aligned + 8192);
assert_eq!(ss.calc_staging_off(1), aligned + 8192 + 4096);
}
#[test]
#[should_panic(expected = "too small")]
fn segment_sum_from_slice_too_small() {
SegmentSum::from_slice(&[0u8; 10]);
}
#[test]
fn segment_sum_write_to_oversized_buf() {
let ndatablk = 2;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0, &inode);
let mut buf = vec![0xFFu8; ss_bytes * 2]; ss.write_to(&mut buf);
assert!(buf[ss_bytes..].iter().all(|&b| b == 0xFF));
let restored = SegmentSum::from_slice(&buf[..ss_bytes]);
assert_eq!(restored.blocks.len(), ndatablk);
}
#[test]
#[should_panic(expected = "too small")]
fn segment_sum_write_to_too_small_for_entries() {
let ndatablk = 3;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let _ = ss.update(0, &inode);
let hdrsz = std::mem::size_of::<SegmentHeader>();
let mut buf = vec![0u8; hdrsz]; ss.write_to(&mut buf);
}
#[test]
fn segment_sum_many_blocks_crossing_4kib() {
let ndatablk = 500;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0, &inode);
assert!(ss_bytes > 4096, "test should span multiple 4KiB pages");
let mut buf = vec![0u8; ss_bytes];
ss.write_to(&mut buf);
let restored = SegmentSum::from_slice(&buf);
assert_eq!(restored.blocks.len(), ndatablk);
for i in 0..ndatablk {
assert_eq!(restored.blocks[i].blkidx, i as u64);
assert_eq!(restored.blocks[i].blkptr, 0x4000_0001_0000_0000 + i as u64);
}
}
#[test]
fn segment_sum_write_is_idempotent() {
let ndatablk = 10;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0xCAFE, &inode);
let mut buf1 = vec![0u8; ss_bytes];
let mut buf2 = vec![0u8; ss_bytes];
ss.write_to(&mut buf1);
ss.write_to(&mut buf2);
assert_eq!(buf1, buf2, "SegmentSum::write_to is not byte-identical on repeat");
}
#[test]
fn segment_sum_round_trip_reserialize_matches() {
let ndatablk = 5;
let mut ss = make_segment_sum(ndatablk);
let inode = InodeRaw::default();
let ss_bytes = ss.update(0xBEEF, &inode);
let mut buf1 = vec![0u8; ss_bytes];
ss.write_to(&mut buf1);
let ss2 = SegmentSum::from_slice(&buf1);
let mut buf2 = vec![0u8; ss_bytes];
ss2.write_to(&mut buf2);
assert_eq!(buf1, buf2, "round-trip write-read-write mismatch");
}
#[test]
fn segid_to_staging_file_id_format() {
assert_eq!(Segment::segid_to_staging_file_id(SegmentId::new(1)), "0000000001");
assert_eq!(Segment::segid_to_staging_file_id(SegmentId::new(u32::MAX)), "4294967295");
assert_eq!(Segment::segid_to_staging_file_id(SegmentId::with_part(1, 0)), "0000000001.0");
assert_eq!(Segment::segid_to_staging_file_id(SegmentId::with_part(1, 7)), "0000000001.7");
assert_ne!(Segment::segid_to_staging_file_id(SegmentId::with_part(1, 0)),
Segment::segid_to_staging_file_id(SegmentId::new(1)));
}
#[test]
fn writer_calc_ss_aligned_bytes() {
let hdr_sz = std::mem::size_of::<SegmentHeader>();
let entry_sz = std::mem::size_of::<SegmentBlockEntryRaw>();
let raw = hdr_sz;
let expected = (raw + 4095) >> 12 << 12;
assert_eq!(expected, 4096);
let n = (4096 - hdr_sz) / entry_sz + 1;
let raw = hdr_sz + n * entry_sz;
let expected = (raw + 4095) >> 12 << 12;
assert!(expected > 4096);
assert_eq!(expected % 4096, 0);
}
}