use std::fs::{File, OpenOptions};
use std::io::{self, BufWriter, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::time::Instant;
use kevy_resp::ArgvView;
use kevy_store::Store;
use crate::{estimate_multibulk_bytes, write_multibulk};
pub const AOF_MAGIC: &[u8; 9] = b"KEVYAOF1\n";
pub(crate) const AOF_BUF_CAP: usize = 256 * 1024;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Fsync {
Always,
EverySec,
No,
}
pub struct Aof {
pub(crate) file: BufWriter<File>,
pub(crate) in_txn: bool,
pub(crate) path: PathBuf,
pub(crate) fsync: Fsync,
pub(crate) dirty: bool,
pub(crate) last_sync: Instant,
pub(crate) size_bytes: u64,
pub(crate) size_at_last_rewrite: u64,
pub(crate) rewrites_total: u64,
pub(crate) deferred: bool,
pub(crate) rewrite_tee: Option<Vec<u8>>,
pub(crate) tee_spare: Option<Vec<u8>>,
pub(crate) swap_trash: Option<PathBuf>,
pub(crate) swap_hold: bool,
open_quarantine: Option<PathBuf>,
pub(crate) last_rewrite_at: Instant,
pub(crate) format: crate::AofFormat,
scratch: Vec<u8>,
pub(crate) queue: Option<Vec<u8>>,
pub(crate) queued_offset: u64,
pub(crate) queued_seq: u64,
}
pub struct RewritePlan {
pub body: Vec<u8>,
pub tmp: PathBuf,
pub keys: u64,
}
#[derive(Debug, Clone, Copy)]
pub struct RewriteStats {
pub keys: u64,
pub bytes: u64,
}
impl Aof {
#[must_use]
pub fn format(&self) -> crate::AofFormat {
self.format
}
pub fn open(path: &Path, fsync: Fsync) -> io::Result<Self> {
Self::open_with_repair(path, fsync, false)
}
pub fn open_with_repair(path: &Path, fsync: Fsync, resync: bool) -> io::Result<Self> {
let mut file = OpenOptions::new().create(true).append(true).open(path)?;
let mut size = file.metadata().map_or(0, |m| m.len());
let mut quarantined = None;
let mut format = crate::AofFormat::V2;
if size == 0 {
file.write_all(crate::record::AOF2_MAGIC)?;
file.sync_data()?;
size = crate::record::AOF2_MAGIC.len() as u64;
} else {
format = crate::replay::sniff_format(path)?;
quarantined = crate::aof_util::repair_tail(path, &mut file, &mut size, resync)?;
}
Ok(Aof {
in_txn: false,
file: BufWriter::with_capacity(AOF_BUF_CAP, file),
path: path.to_path_buf(),
fsync,
dirty: false,
last_sync: Instant::now(),
size_bytes: size,
size_at_last_rewrite: size,
rewrites_total: 0,
deferred: false,
rewrite_tee: None,
tee_spare: None,
swap_trash: None,
swap_hold: false,
open_quarantine: quarantined,
last_rewrite_at: Instant::now(),
format,
scratch: Vec::new(),
queue: None,
queued_offset: size,
queued_seq: 0,
})
}
#[inline]
pub fn open_quarantine(&self) -> Option<&Path> {
self.open_quarantine.as_deref()
}
#[inline]
pub(crate) fn last_rewrite_at(&self) -> Instant {
self.last_rewrite_at
}
#[inline]
pub fn fsync_policy(&self) -> Fsync {
self.fsync
}
pub fn set_fsync(&mut self, fsync: Fsync) -> io::Result<()> {
let upgrading_to_always =
matches!(fsync, Fsync::Always) && !matches!(self.fsync, Fsync::Always);
self.fsync = fsync;
if upgrading_to_always {
self.flush_queued()?;
}
if upgrading_to_always && self.dirty {
self.file.flush()?;
self.file.get_ref().sync_data()?;
self.dirty = false;
self.last_sync = Instant::now();
}
Ok(())
}
fn write_scratch_to_file(&mut self) -> io::Result<()> {
match self.format {
crate::AofFormat::V2 => {
self.file.write_all(&(self.scratch.len() as u32).to_le_bytes())?;
self.file.write_all(&crate::crc32c::crc32c(&self.scratch).to_le_bytes())?;
self.file.write_all(&self.scratch)
}
crate::AofFormat::V1 => self.file.write_all(&self.scratch),
}
}
pub fn append<A: ArgvView + ?Sized>(&mut self, args: &A) -> io::Result<()> {
self.scratch.clear();
write_multibulk(&mut self.scratch, args)?;
if let Some(q) = &mut self.queue {
match self.format {
crate::AofFormat::V2 => {
q.extend_from_slice(&(self.scratch.len() as u32).to_le_bytes());
q.extend_from_slice(&crate::crc32c::crc32c(&self.scratch).to_le_bytes());
q.extend_from_slice(&self.scratch);
}
crate::AofFormat::V1 => q.extend_from_slice(&self.scratch),
}
self.queued_seq += 1;
} else {
self.write_scratch_to_file()?;
}
if let Some(tee) = &mut self.rewrite_tee {
crate::record::write_record(tee, &self.scratch)?;
}
let overhead = match self.format {
crate::AofFormat::V2 => crate::record::RECORD_HEADER as u64,
crate::AofFormat::V1 => 0,
};
self.size_bytes =
self.size_bytes.saturating_add(estimate_multibulk_bytes(args)).saturating_add(overhead);
match self.fsync {
Fsync::Always if self.deferred || self.queue.is_some() => self.dirty = true,
Fsync::Always => {
self.file.flush()?;
self.file.get_ref().sync_data()?;
}
Fsync::EverySec | Fsync::No => self.dirty = true,
}
Ok(())
}
pub fn truncate(&mut self) -> io::Result<()> {
self.flush_queued()?;
self.file.flush()?;
let f = self.file.get_mut();
f.set_len(0)?;
f.seek(SeekFrom::Start(0))?; f.write_all(crate::record::AOF2_MAGIC)?;
f.sync_all()?;
self.dirty = false;
self.format = crate::AofFormat::V2; self.size_bytes = crate::record::AOF2_MAGIC.len() as u64;
self.queued_offset = self.size_bytes;
self.size_at_last_rewrite = crate::record::AOF2_MAGIC.len() as u64;
self.last_rewrite_at = Instant::now();
Ok(())
}
#[inline]
pub fn size_bytes(&self) -> u64 {
self.size_bytes
}
#[inline]
pub fn size_at_last_rewrite(&self) -> u64 {
self.size_at_last_rewrite
}
pub fn anchor_rewrite_baseline(&mut self, bytes: u64) {
if !self.is_rewriting() {
self.size_at_last_rewrite = bytes.max(crate::record::AOF2_MAGIC.len() as u64);
}
}
#[inline]
pub fn rewrites_total(&self) -> u64 {
self.rewrites_total
}
pub fn rewrite_from(&mut self, store: &Store) -> io::Result<RewriteStats> {
self.flush_queued()?;
self.file.flush()?;
let tmp = crate::aof_util::rewrite_tmp_path(&self.path);
let (keys, bytes) = crate::dump_aof(&tmp, store)?;
std::fs::rename(&tmp, &self.path)?;
let f = OpenOptions::new().append(true).open(&self.path)?;
self.file = BufWriter::with_capacity(AOF_BUF_CAP, f);
self.format = crate::AofFormat::V2; self.size_bytes = bytes;
self.queued_offset = bytes;
self.size_at_last_rewrite = bytes;
self.last_rewrite_at = Instant::now();
self.dirty = false;
self.rewrites_total = self.rewrites_total.saturating_add(1);
Ok(RewriteStats { keys, bytes })
}
#[inline]
pub fn is_rewriting(&self) -> bool {
self.rewrite_tee.is_some()
}
}