use std::fs::{File, OpenOptions};
use std::io::{self, BufWriter, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use kevy_resp::ArgvView;
use kevy_store::Store;
use crate::{
dump_store_to_buf, estimate_multibulk_bytes, write_multibulk,
};
pub const AOF_MAGIC: &[u8; 9] = b"KEVYAOF1\n";
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,
path: PathBuf,
pub(crate) fsync: Fsync,
pub(crate) dirty: bool,
pub(crate) last_sync: Instant,
size_bytes: u64,
size_at_last_rewrite: u64,
rewrites_total: u64,
pub(crate) deferred: bool,
rewrite_tee: Option<Vec<u8>>,
open_quarantine: Option<PathBuf>,
last_rewrite_at: Instant,
pub(crate) format: crate::AofFormat,
scratch: Vec<u8>,
}
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,
open_quarantine: quarantined,
last_rewrite_at: Instant::now(),
format,
scratch: Vec::new(),
})
}
#[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.dirty {
self.file.flush()?;
self.file.get_ref().sync_data()?;
self.dirty = false;
self.last_sync = Instant::now();
}
Ok(())
}
pub fn append<A: ArgvView + ?Sized>(&mut self, args: &A) -> io::Result<()> {
self.scratch.clear();
write_multibulk(&mut self.scratch, args)?;
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)?,
}
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.dirty = true,
Fsync::Always => {
self.file.flush()?;
self.file.get_ref().sync_data()?;
}
Fsync::EverySec | Fsync::No => self.dirty = true,
}
Ok(())
}
pub fn sync_now(&mut self) -> io::Result<()> {
if self.dirty {
self.file.flush()?;
self.file.get_ref().sync_data()?;
self.dirty = false;
self.last_sync = Instant::now();
}
Ok(())
}
pub fn maybe_sync(&mut self) -> io::Result<()> {
if matches!(self.fsync, Fsync::EverySec)
&& self.dirty
&& self.last_sync.elapsed() >= Duration::from_secs(1)
{
self.file.flush()?;
self.file.get_ref().sync_data()?;
self.dirty = false;
self.last_sync = Instant::now();
}
Ok(())
}
pub fn truncate(&mut self) -> io::Result<()> {
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.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
}
#[inline]
pub fn rewrites_total(&self) -> u64 {
self.rewrites_total
}
pub fn rewrite_from(&mut self, store: &Store) -> io::Result<RewriteStats> {
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.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()
}
pub fn begin_concurrent_rewrite(&mut self, store: &Store) -> io::Result<RewritePlan> {
let (body, keys) = dump_store_to_buf(store, crate::AofFormat::V2);
self.rewrite_tee = Some(Vec::new());
Ok(RewritePlan {
body,
tmp: crate::aof_util::rewrite_tmp_path(&self.path),
keys,
})
}
pub fn finish_concurrent_rewrite(&mut self, tmp: &Path, keys: u64) -> io::Result<RewriteStats> {
let tee = self.rewrite_tee.take().unwrap_or_default();
{
let mut f = OpenOptions::new().append(true).open(tmp)?;
f.write_all(&tee)?;
f.sync_all()?;
}
std::fs::rename(tmp, &self.path)?;
let f = OpenOptions::new().append(true).open(&self.path)?;
let bytes = f.metadata().map_or(0, |m| m.len());
self.file = BufWriter::with_capacity(AOF_BUF_CAP, f);
self.format = crate::AofFormat::V2; self.size_bytes = 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 })
}
pub fn abort_concurrent_rewrite(&mut self) {
self.rewrite_tee = None;
}
pub fn begin_view_rewrite(&mut self) -> io::Result<std::path::PathBuf> {
self.file.flush()?;
self.rewrite_tee = Some(Vec::new());
Ok(crate::aof_util::rewrite_tmp_path(&self.path))
}
}