roreo_common/

manifest.rs

use std::{
  fs::File,
  io::{Read, Seek, SeekFrom, Write},
  mem::MaybeUninit,
  path::PathBuf, sync::Arc,
};

pub use crate::pb::{ManifestChange as PbManifestChange, ManifestChangeSet};
use indexmap::IndexMap;
use parking_lot::Mutex;
use prost::Message;

const MANIFEST_FILENAME: &str = "MANIFEST";
const MANIFEST_REWRITE_FILENAME: &str = "MANIFEST-REWRITE";

// The magic version number. It is allocated 2 bytes, so it's value must be <= u16::MAX
const MAGIC_VERSION: u16 = 0;
// Has to be 4 bytes. The value can nerver change, ever, anyway.
const MAGIC_TEXT: [u8; 4] = *b"al8n";

const DEFAULT_DELETIONS_REWRITE_THRESHOLD: u64 = 10_000;
const DEFAULT_DELETIONS_RATIO: u64 = 10;

/// Manifest file errors.
#[derive(Debug, thiserror::Error)]
pub enum ManifestError {
  /// IO error
  #[error("roreo: manifest io: {0}")]
  IO(#[from] std::io::Error),
  /// manifest removes non-existing 
  #[error("roreo: manifest removes non-existing {0} file {1}")]
  RemoveNotExists(FileType, u64),
  /// manifest adds already existing
  #[error("roreo: manifest adds already existing {0} file {1}")]
  AddExists(FileType, u64),
  /// manifest file not exists
  #[error("roreo: manifest file not exists, required for read-only db")]
  NotFound,
  /// manifest has bad magic
  #[error("roreo: manifest has bad magic")]
  BadMagic,
  /// manifest checksum mismatch
  #[error("roreo: manifest has checksum mismatch")]
  ChecksumMismatch,
  /// manifest cannot open DB because the external magic number doesn't match
  #[error("roreo: manifest cannot open DB because the external magic number doesn't match (expected {0} != {1}).")]
  ExtrenalMagicMismtach(u16, u16),
  /// manifest file maybe corrupted
  #[error("roreo: manifest buffer length {buf_len} greater than file size {file_size}, file may be corrupted")]
  Corrupted {
    /// the length of the buffer
    buf_len: u64,
    /// the size of the file
    file_size: u64
  },
  /// manifest decode error
  #[error("roreo: manifest decode error: {0}")]
  DecodeError(#[from] prost::DecodeError),
  /// manifest unkonwn operation type
  #[error("roreo: manifest unknown manifest operation, manifest file may be corrupted")]
  InvalidOperation,
  /// manifest unkonwn file type
  #[error("roreo: manifest unknown file type, manifest file may be corrupted")]
  InvalidFileType,
  /// manifest unsupported version
  #[error("roreo: manifest unsupported version {version} (we support {supported_version}).\n please see https://github.com/al8n/roreo/blob/main/sorted/version.md on how to fix this.")]
  UnsupportedVersion {
    /// the version used in the application
    version: u16,
    /// the version supported by the roreo
    supported_version: u16,
  },
}

/// The type of a file in the manifest.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum FileType {
  /// WAL file
  WAL,
  /// Data file
  Data,
  /// Hint file
  Hint,
}

impl FileType {
  #[inline]
  const fn wal() -> Self {
    Self::WAL
  }

  #[inline]
  const fn data() -> Self {
    Self::Data
  }

  #[inline]
  const fn hint() -> Self {
    Self::Hint
  }
}

impl core::fmt::Display for FileType {
  fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
    match self {
      Self::WAL => write!(f, "wal"),
      Self::Data => write!(f, "data"),
      Self::Hint => write!(f, "hint"),
    }
  }
}

/// The operation of a file in the manifest.
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[repr(u8)]
pub enum FileOp {
  /// Create a file
  Create,
  /// Delete a file
  Delete,
}


#[viewit::viewit(
  vis_all = "pub(crate)",
  setters(vis_all = "pub"),
  getters(vis_all = "pub")
)]
#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
/// The in memory version atomic manifest change.
pub struct ManifestChange {
  // the fields order is the order of encoding in the manifest file
  file_op: FileOp,
  file_type: FileType,
  compression: u8,
  encryption_algo: u8,
  fid: u64,
}

impl From<ManifestChange> for PbManifestChange {
  fn from(value: ManifestChange) -> Self {
    Self {
      meta: u32::from_le_bytes([
        value.file_op as u8,
        value.file_type as u8,
        value.compression,
        value.encryption_algo,
      ]),
      fid: value.fid,
    }
  }
}

impl TryFrom<PbManifestChange> for ManifestChange {
  type Error = ManifestError;

  fn try_from(value: PbManifestChange) -> Result<Self, Self::Error> {
    let meta = value.meta.to_le_bytes();
    Ok(Self {
      file_op: match meta[0] {
        0 => FileOp::Create,
        1 => FileOp::Delete,
        _ => return Err(ManifestError::InvalidOperation),
      },
      file_type: match meta[1] {
        0 => FileType::wal(),
        1 => FileType::data(),
        2 => FileType::hint(),
        _ => return Err(ManifestError::InvalidFileType),
      },
      compression: meta[2],
      encryption_algo: meta[3],
      fid: value.fid,
    })
  }
}

/// The options used for the manifest file.
#[viewit::viewit(
  vis_all = "pub(crate)",
  setters(vis_all = "pub"),
  getters(vis_all = "pub")
)]
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ManifestOptions {
  #[viewit(getter(style = "ref"))]
  dir: PathBuf,
  deletion_threshold: u64,
  deletion_ratio: u64,
  external_magic_version: u16,
  read_only: bool,
}

impl ManifestOptions {
  /// Create a default ManifestOptions with the given directory.
  #[inline]
  pub const fn new(dir: PathBuf) -> Self {
    Self {
      dir,
      deletion_threshold: DEFAULT_DELETIONS_REWRITE_THRESHOLD,
      deletion_ratio: DEFAULT_DELETIONS_RATIO,
      external_magic_version: 0,
      read_only: false,
    }
  }
}

#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct ManifestInfo {
  compression: u8,
  encryption_algo: u8,
}

#[derive(Debug, Default)]
#[repr(transparent)]
struct WALManifest {
  files: IndexMap<u64, ManifestInfo>,
}

#[derive(Debug, Default)]
#[repr(transparent)]
struct DataManifest {
  files: IndexMap<u64, ManifestInfo>,
}

#[derive(Debug, Default)]
#[repr(transparent)]
struct HintManifest {
  files: IndexMap<u64, ManifestInfo>,
}

#[derive(Debug, Default)]
struct Manifest {
  wal: WALManifest,
  data: DataManifest,
  hint: HintManifest,

  // Contains total number of creation and deletion changes in the manifest -- used to compute
  // whether it'd be useful to rewrite the manifest.
  creations: u64,
  deletions: u64,
}

impl From<&Manifest> for Vec<PbManifestChange> {
  fn from(value: &Manifest) -> Self {
    let mut changes =
      Vec::with_capacity(value.wal.files.len() + value.data.files.len() + value.hint.files.len());

    for (&fid, info) in &value.wal.files {
      changes.push(PbManifestChange {
        meta: u32::from_le_bytes([
          FileOp::Create as u8,
          FileType::WAL as u8,
          info.compression,
          info.encryption_algo,
        ]),
        fid,
      });
    }

    for (&fid, info) in &value.data.files {
      changes.push(PbManifestChange {
        meta: u32::from_le_bytes([
          FileOp::Create as u8,
          FileType::Data as u8,
          info.compression,
          info.encryption_algo,
        ]),
        fid,
      });
    }

    for (&fid, info) in &value.hint.files {
      changes.push(PbManifestChange {
        meta: u32::from_le_bytes([
          FileOp::Create as u8,
          FileType::Hint as u8,
          info.compression,
          info.encryption_algo,
        ]),
        fid,
      });
    }

    changes
  }
}

macro_rules! insert_bail {
  ($manifest:ident::$change:ident::$name:ident: $name_str: literal) => {
    if $manifest
      .$name
      .files
      .insert(
        $change.fid,
        ManifestInfo {
          compression: $change.compression,
          encryption_algo: $change.encryption_algo,
        },
      )
      .is_some()
    {
      tracing::error!(
        target = "manifest",
        concat!(
          "roreo: try to add ",
          $name_str,
          " file {}, but this fid already exists in manifest"
        ),
        $change.fid
      );
      return Err(ManifestError::AddExists(FileType::$name(), $change.fid));
    } else {
      tracing::info!(
        target = "manifest",
        concat!("roreo: add ", $name_str, " file {} to manifest"),
        $change.fid
      );
    }
  };
}

macro_rules! remove_bail {
  ($manifest:ident::$change:ident::$name:ident: $name_str: literal) => {
    match $manifest.$name.files.remove(&$change.fid) {
      Some(_) => {
        tracing::info!(
          target = "manifest",
          concat!("roreo: remove ", $name_str, " file {} from manifest"),
          $change.fid
        );
      }
      None => {
        tracing::error!(
          target = "manifest",
          concat!(
            "roreo: try to remove ",
            $name_str,
            "file {}, but not found in manifest"
          ),
          $change.fid
        );
        return Err(ManifestError::RemoveNotExists(
          FileType::$name(),
          $change.fid,
        ));
      }
    }
  };
}

impl Manifest {
  #[inline]
  fn apply_changes(
    &mut self,
    changes: Vec<ManifestChange>,
  ) -> Result<ManifestChangeSet, ManifestError> {
    let mut creations = 0;
    let mut deletions = 0;
    let set = ManifestChangeSet {
      changes: changes
        .into_iter()
        .map(|change| {
          match (change.file_type, change.file_op) {
            (FileType::WAL, FileOp::Create) => {
              insert_bail!(self::change::wal: "wal");
              creations += 1;
            }
            (FileType::WAL, FileOp::Delete) => {
              remove_bail!(self::change::wal: "wal");
              deletions += 1;
            }
            (FileType::Data, FileOp::Create) => {
              insert_bail!(self::change::data: "data");
              creations += 1;
            }
            (FileType::Data, FileOp::Delete) => {
              remove_bail!(self::change::data: "data");
              deletions += 1;
            }
            (FileType::Hint, FileOp::Create) => {
              insert_bail!(self::change::hint: "hint");
              creations += 1;
            }
            (FileType::Hint, FileOp::Delete) => {
              remove_bail!(self::change::hint: "hint");
              deletions += 1;
            }
          }
          Ok(change.into())
        })
        .collect::<Result<Vec<_>, ManifestError>>()?,
    };

    self.deletions += deletions;
    self.creations += creations;
    Ok(set)
  }

  #[inline]
  fn apply_change_set(&mut self, changes: ManifestChangeSet) -> Result<(), ManifestError> {
    for change in changes.changes {
      let change: ManifestChange = change.try_into()?;
      match (change.file_type, change.file_op) {
        (FileType::WAL, FileOp::Create) => {
          insert_bail!(self::change::wal: "wal");
          self.creations += 1;
        }
        (FileType::WAL, FileOp::Delete) => {
          remove_bail!(self::change::wal: "wal");
          self.deletions += 1;
        }
        (FileType::Data, FileOp::Create) => {
          insert_bail!(self::change::data: "data");
          self.creations += 1;
        }
        (FileType::Data, FileOp::Delete) => {
          remove_bail!(self::change::data: "data");
          self.deletions += 1;
        }
        (FileType::Hint, FileOp::Create) => {
          insert_bail!(self::change::hint: "hint");
          self.creations += 1;
        }
        (FileType::Hint, FileOp::Delete) => {
          remove_bail!(self::change::hint: "hint");
          self.deletions += 1;
        }
      }
    }
    Ok(())
  }
}

#[derive(Debug)]
struct ManifestFileInner {
  manifest: Manifest,
  wal: MaybeUninit<File>,
}

impl ManifestFileInner {
  fn rewrite(&mut self, manifest_opts: &ManifestOptions) -> Result<u64, ManifestError> {
    // close the old manifest file
    // Safety: we have a mutex to ensure that only one thread can write to the manifest at a time, and we also
    unsafe { self.wal.assume_init_drop() };

    let changes = Vec::<PbManifestChange>::from(&self.manifest);
    let creations = changes.len() as u64;
    let dir = &manifest_opts.dir;
    let rewrite_file_name = dir.join(MANIFEST_REWRITE_FILENAME);
    let mut opts = std::fs::OpenOptions::new();
    let mut file = opts
      .create(true)
      .write(true)
      .truncate(true)
      .read(true)
      .open(&rewrite_file_name)?;

    // magic bytes are structured as
    // +---------------------+-------------------------+-----------------------+
    // | magicText (4 bytes) | externalMagic (2 bytes) | roreoMagic (2 bytes) |
    // +---------------------+-------------------------+-----------------------+
    let set = ManifestChangeSet { changes }.encode_to_vec();
    let mut buf_writter = std::io::BufWriter::with_capacity(
      4 // magic text
      + 2 // external magic
      + 2 // roreo magic
      + 4 // data length
      + 4 // cks
      + set.len(),
      &mut file,
    );

    buf_writter.write_all(&MAGIC_TEXT).unwrap();
    buf_writter
      .write_all(&manifest_opts.external_magic_version.to_le_bytes())
      .unwrap();
    buf_writter.write_all(&MAGIC_VERSION.to_le_bytes()).unwrap();
    buf_writter
      .write_all(&(set.len() as u32).to_le_bytes())
      .unwrap();
    buf_writter
      .write_all(crc32fast::hash(&set).to_le_bytes().as_ref())
      .unwrap();
    buf_writter.write_all(&set).unwrap();
    buf_writter.flush()?;

    // close the rewrite file for renaming
    drop(buf_writter);
    drop(file);

    let manifest_file_path = dir.join(MANIFEST_FILENAME);
    std::fs::rename(&rewrite_file_name, &manifest_file_path)?;

    opts
      .create(false)
      .write(true)
      .append(true)
      .read(true)
      .open(&manifest_file_path)
      .and_then(|mut f| f.seek(SeekFrom::End(0)).map(|_| self.wal.write(f)))
      .and_then(|_| File::open(&manifest_opts.dir).and_then(|dir| dir.sync_all()))
      .map(|_| creations)
      .map_err(From::from)
  }
}

/// The write-ahead log manifest file
#[derive(Debug, Clone)]
pub struct ManifestFile {
  opts: ManifestOptions,
  inner: Arc<Mutex<ManifestFileInner>>,
}

impl ManifestFile {
  /// Open or create a manifest file
  pub fn open(opt: ManifestOptions) -> Result<Self, ManifestError> {
    let file_path = opt.dir.join(MANIFEST_FILENAME);
    let mut opts = std::fs::OpenOptions::new();
    if opt.read_only {
      opts.read(true);
    } else {
      opts.read(true).write(true).create(true);
    }

    if !file_path.exists() {
      if opt.read_only {
        return Err(ManifestError::NotFound);
      }
      return opts
        .open(&file_path)
        .and_then(|mut f| {
          let mut magic_buf = [0u8; 8];
          magic_buf[..4].copy_from_slice(&MAGIC_TEXT);
          magic_buf[4..6].copy_from_slice(&opt.external_magic_version.to_le_bytes());
          magic_buf[6..].copy_from_slice(&MAGIC_VERSION.to_le_bytes());
          f.write_all(&magic_buf).map(|_| Self {
            opts: opt,
            inner: Arc::new(Mutex::new(ManifestFileInner {
              manifest: Default::default(),
              wal: MaybeUninit::new(f),
            })),
          })
        })
        .map_err(From::from);
    }

    opts.open(&file_path).map_err(From::from).and_then(|mut f| {
      replay(&mut f, opt.external_magic_version, opt.read_only).map(|m| Self {
        opts: opt,
        inner: Arc::new(Mutex::new(ManifestFileInner {
          manifest: m,
          wal: MaybeUninit::new(f),
        })),
      })
    })
  }

  /// Write a change to the manifest file, this fn is atomic
  pub fn write_change(&self, change: ManifestChange) -> Result<(), ManifestError> {
    let mut inner = self.inner.lock();
    if inner.manifest.deletions > self.opts.deletion_threshold
      && inner.manifest.deletions
        > self.opts.deletion_ratio * (inner.manifest.creations + inner.manifest.deletions)
    {
      inner.rewrite(&self.opts).map(|creations| {
        inner.manifest.creations = creations;
        inner.manifest.deletions = 0;
      })
    } else {
      let encoded = ManifestChangeSet {
        changes: vec![PbManifestChange::from(change)],
      }
      .encode_to_vec();
      let encoded_len = encoded.len();
      let mut buf = std::io::BufWriter::with_capacity(
        4 + 4 + encoded_len,
        // Safety: We know the wal is initialized.
        unsafe { inner.wal.assume_init_mut() },
      );

      // Panic: impossible to fail, because the buffer is pre-allocated
      buf.write_all(&(encoded_len as u32).to_le_bytes()).unwrap();
      buf
        .write_all(crc32fast::hash(&encoded).to_le_bytes().as_ref())
        .unwrap();
      buf.write_all(&encoded).unwrap();
      buf.flush().map_err(From::from)
    }
  }

  /// Write a list of changes to the manifest file, this fn is atomic
  pub fn write_changes(&self, changes: Vec<ManifestChange>) -> Result<(), ManifestError> {
    let mut inner = self.inner.lock();
    let change_set = inner.manifest.apply_changes(changes)?;

    // Rewrite manifest if it'd shrink by 1/10 and it's big enough to care
    if inner.manifest.deletions > self.opts.deletion_threshold
      && inner.manifest.deletions
        > self.opts.deletion_ratio * (inner.manifest.creations + inner.manifest.deletions)
    {
      inner.rewrite(&self.opts).map(|creations| {
        inner.manifest.creations = creations;
        inner.manifest.deletions = 0;
      })
    } else {
      let encoded = change_set.encode_to_vec();
      let encoded_len = encoded.len();
      let mut buf = std::io::BufWriter::with_capacity(
        4 + 4 + encoded_len,
        // Safety: We know the wal is initialized.
        unsafe { inner.wal.assume_init_mut() },
      );

      // Panic: impossible to fail, because the buffer is pre-allocated
      buf.write_all(&(encoded_len as u32).to_le_bytes()).unwrap();
      buf
        .write_all(crc32fast::hash(&encoded).to_le_bytes().as_ref())
        .unwrap();
      buf.write_all(&encoded).unwrap();
      buf.flush().map_err(From::from)
    }
  }
}

#[inline]
fn replay(
  file: &mut File,
  external_magic: u16,
  read_only: bool,
) -> Result<Manifest, ManifestError> {
  let stat = file.metadata()?;
  let mut reader = std::io::BufReader::new(file);

  let mut magic_buf = [0; 8];
  reader.read_exact(&mut magic_buf).map_err(|e| {
    tracing::error!(target = "manifest", "roreo: read magic failed: {}", e);
    ManifestError::BadMagic
  })?;

  if magic_buf[..4] != MAGIC_TEXT {
    tracing::error!(target = "manifest", "roreo: magic mismatch");
    return Err(ManifestError::BadMagic);
  }

  let ext_version = u16::from_le_bytes(magic_buf[4..6].try_into().unwrap());
  let version = u16::from_le_bytes(magic_buf[6..8].try_into().unwrap());

  if version != MAGIC_VERSION {
    return Err(ManifestError::UnsupportedVersion {
      version,
      supported_version: MAGIC_VERSION,
    });
  }

  if ext_version != external_magic {
    tracing::error!(
      target = "manifest",
      "roreo: external magic mismatch expected {}, found {}",
      external_magic,
      ext_version
    );
    return Err(ManifestError::ExtrenalMagicMismtach(
      external_magic,
      ext_version,
    ));
  }

  let mut manifest = Manifest::default();

  let mut offset = 8u64;
  loop {
    let mut len_crc_buf = [0; 8];
    if let Err(e) = reader.read_exact(&mut len_crc_buf) {
      if e.kind() == std::io::ErrorKind::UnexpectedEof {
        break;
      }
      return Err(e.into());
    }
    offset += 8;

    let len = u32::from_le_bytes(len_crc_buf[..4].try_into().unwrap()) as u64;
    // Sanity check to ensure we don't over-allocate memory.
    if len > stat.len() {
      return Err(ManifestError::Corrupted {
        buf_len: len,
        file_size: stat.len(),
      });
    }

    let mut buf = vec![0; len as usize];
    if let Err(e) = reader.read_exact(&mut buf) {
      if e.kind() == std::io::ErrorKind::UnexpectedEof {
        break;
      }
      return Err(e.into());
    }
    offset += len;

    let crc = u32::from_le_bytes(len_crc_buf[4..8].try_into().unwrap());
    if crc != crc32fast::hash(&buf) {
      return Err(ManifestError::ChecksumMismatch);
    }

    ManifestChangeSet::decode(buf.as_slice())
      .map_err(From::from)
      .and_then(|cs| manifest.apply_change_set(cs))?;
  }
  let file = reader.into_inner();
  if !read_only {
    // Truncate file so we don't have a half-written entry at the end.
    file.set_len(offset)?;
  }
  file.seek(SeekFrom::End(0))?;

  Ok(manifest)
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn test_manifest_basic() {
    let dir = tempfile::tempdir().unwrap();
    let opts = ManifestOptions::new(dir.path().to_path_buf()).set_deletion_threshold(100);
    let manifest = ManifestFile::open(opts.clone()).unwrap();

    let mut changes = vec![];
    for i in 0..1000usize {
      changes.push(vec![ManifestChange {
        file_op: FileOp::Create,
        file_type: FileType::WAL,
        compression: 0,
        encryption_algo: 0,
        fid: i as u64,
      }]);
      changes.push(vec![
        ManifestChange {
          file_op: FileOp::Create,
          file_type: FileType::Data,
          compression: 0,
          encryption_algo: 0,
          fid: i as u64,
        },
        ManifestChange {
          file_op: FileOp::Create,
          file_type: FileType::Hint,
          compression: 0,
          encryption_algo: 0,
          fid: i as u64,
        },
      ]);
      changes.push(vec![ManifestChange {
        file_op: FileOp::Delete,
        file_type: FileType::WAL,
        compression: 0,
        encryption_algo: 0,
        fid: i as u64,
      }]);
    }

    for change in changes {
      manifest.write_changes(change).unwrap();
    }
    drop(manifest);

    let manifest = ManifestFile::open(opts).unwrap();

    let inner = manifest.inner.lock();
    assert!(inner.manifest.wal.files.is_empty());
    assert_eq!(inner.manifest.data.files.len(), 1000);
    assert_eq!(inner.manifest.hint.files.len(), 1000);
    assert_eq!(
      inner.manifest.data.files,
      (0..1000)
        .map(|idx| (
          idx,
          ManifestInfo {
            compression: 0,
            encryption_algo: 0
          }
        ))
        .collect::<IndexMap<_, _>>()
    );
    assert_eq!(
      inner.manifest.hint.files,
      (0..1000)
        .map(|idx| (
          idx,
          ManifestInfo {
            compression: 0,
            encryption_algo: 0
          }
        ))
        .collect::<IndexMap<_, _>>()
    );
  }
}