use std::sync::{Arc, Mutex};
use pagesift_store::{Store, Tasks};
use super::store::{list_key, meta_key};
pub trait JobLog: Send + Sync {
fn append(&self, op: Vec<u8>) -> Result<(), String>;
fn max_op(&self) -> usize {
usize::MAX
}
}
pub struct JobMirror {
log: Arc<dyn JobLog>,
lock: Mutex<()>,
failures: std::sync::atomic::AtomicU64,
}
const KV_PUT: u8 = 1;
const KV_DEL: u8 = 2; const LIST_PUSH: u8 = 3; const LIST_CLEAR: u8 = 4; const SET_INSERT: u8 = 5; const SET_CLEAR: u8 = 6; const LIST_POP: u8 = 7;
const LISTS: &[&str] = &["pages", "errors", "frontier", "notes"];
fn put_bytes(out: &mut Vec<u8>, b: &[u8]) {
out.extend_from_slice(&(b.len() as u64).to_le_bytes());
out.extend_from_slice(b);
}
fn op(tag: u8, a: &[u8], n: Option<u64>, b: Option<&[u8]>) -> Vec<u8> {
let mut out = vec![tag];
put_bytes(&mut out, a);
if let Some(n) = n {
out.extend_from_slice(&n.to_le_bytes());
}
if let Some(b) = b {
put_bytes(&mut out, b);
}
out
}
struct Reader<'a>(&'a [u8]);
impl<'a> Reader<'a> {
fn u64(&mut self) -> Option<u64> {
let (h, rest) = self.0.split_first_chunk::<{ size_of::<u64>() }>()?;
self.0 = rest;
Some(u64::from_le_bytes(*h))
}
fn bytes(&mut self) -> Option<&'a [u8]> {
let n = usize::try_from(self.u64()?).ok()?;
let (b, rest) = (self.0.get(..n)?, self.0.get(n..)?);
self.0 = rest;
Some(b)
}
}
pub(crate) struct Writer<'a> {
t: Tasks<'a>,
ops: Option<Vec<Vec<u8>>>,
}
impl Writer<'_> {
pub fn kv_get(&self, key: &str) -> pagesift_store::Result<Option<Vec<u8>>> {
self.t.kv_get(key)
}
pub fn kv_put(&mut self, key: &str, value: &[u8]) -> pagesift_store::Result<()> {
self.t.kv_put(key, value)?;
self.note(|| op(KV_PUT, key.as_bytes(), None, Some(value)));
Ok(())
}
pub fn kv_delete(&mut self, key: &str) -> pagesift_store::Result<()> {
self.t.kv_delete(key)?;
self.note(|| op(KV_DEL, key.as_bytes(), None, None));
Ok(())
}
pub fn list_push(&mut self, list: &str, item: &[u8]) -> pagesift_store::Result<()> {
let i = self.t.list_push(list, item)?;
self.note(|| op(LIST_PUSH, list.as_bytes(), Some(i), Some(item)));
Ok(())
}
pub fn list_clear(&mut self, list: &str) -> pagesift_store::Result<()> {
self.t.list_clear(list)?;
self.note(|| op(LIST_CLEAR, list.as_bytes(), None, None));
Ok(())
}
pub fn set_insert(&mut self, set: &str, member: &[u8]) -> pagesift_store::Result<()> {
self.t.set_insert(set, member)?;
self.note(|| op(SET_INSERT, set.as_bytes(), None, Some(member)));
Ok(())
}
pub fn list_pop_front(&mut self, list: &str) -> pagesift_store::Result<()> {
self.t.list_pop_front(list)?;
self.note(|| op(LIST_POP, list.as_bytes(), None, None));
Ok(())
}
pub fn set_clear(&mut self, set: &str) -> pagesift_store::Result<()> {
self.t.set_clear(set)?;
self.note(|| op(SET_CLEAR, set.as_bytes(), None, None));
Ok(())
}
fn note(&mut self, f: impl FnOnce() -> Vec<u8>) {
if let Some(ops) = &mut self.ops {
ops.push(f());
}
}
}
pub(crate) fn write<R>(
store: &Store,
mirror: Option<&JobMirror>,
f: impl FnOnce(&mut Writer<'_>) -> pagesift_store::Result<R>,
) -> pagesift_store::Result<R> {
let Some(m) = mirror else {
return f(&mut Writer {
t: store.tasks(),
ops: None,
});
};
let _g = m.lock.lock().unwrap_or_else(|e| e.into_inner());
let mut w = Writer {
t: store.tasks(),
ops: Some(Vec::new()),
};
let r = f(&mut w);
let mut failed = None;
for o in w.ops.take().unwrap_or_default() {
if let Err(e) = m.log.append(o) {
m.failures
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
failed = Some(e);
break;
}
}
match (r, failed) {
(Ok(_), Some(e)) => Err(pagesift_store::Error::Io(std::io::Error::other(
Unreplicated(e),
))),
(r, _) => r,
}
}
#[derive(Debug)]
struct Unreplicated(String);
impl std::fmt::Display for Unreplicated {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "replicating to the standby failed: {}", self.0)
}
}
impl std::error::Error for Unreplicated {}
pub(crate) fn unreplicated(e: &pagesift_store::Error) -> bool {
matches!(e, pagesift_store::Error::Io(io) if io.get_ref().is_some_and(|x| x.is::<Unreplicated>()))
}
impl JobMirror {
pub fn new(log: Arc<dyn JobLog>) -> Arc<Self> {
Arc::new(Self {
log,
lock: Mutex::new(()),
failures: std::sync::atomic::AtomicU64::new(0),
})
}
pub fn max_op(&self) -> usize {
self.log.max_op()
}
pub fn failures(&self) -> u64 {
self.failures.load(std::sync::atomic::Ordering::Relaxed)
}
pub fn dump(&self, store: &Store) -> Vec<Vec<u8>> {
let _g = self.lock.lock().unwrap_or_else(|e| e.into_inner());
let t = store.tasks();
let mut out = Vec::new();
for key in t.kv_keys(meta_key("")).unwrap_or_default() {
let Ok(key) = String::from_utf8(key) else {
continue;
};
let id = &key[meta_key("").len()..];
if id.contains('/') {
continue;
}
let Ok(Some(meta)) = t.kv_get(&key) else {
continue;
};
out.push(op(KV_PUT, key.as_bytes(), None, Some(&meta)));
for &what in LISTS {
let list = list_key(id, what);
out.push(op(LIST_CLEAR, list.as_bytes(), None, None));
for (i, item) in t
.list_range(&list, 0, usize::MAX)
.unwrap_or_default()
.into_iter()
.enumerate()
{
out.push(op(LIST_PUSH, list.as_bytes(), Some(i as u64), Some(&item)));
}
}
let hook = crate::webhook::outbox_key(id);
out.push(op(LIST_CLEAR, hook.as_bytes(), None, None));
for (i, item) in t
.list_range(&hook, 0, usize::MAX)
.unwrap_or_default()
.into_iter()
.enumerate()
{
out.push(op(LIST_PUSH, hook.as_bytes(), Some(i as u64), Some(&item)));
}
for k in [crate::webhook::seq_key(id), crate::webhook::ack_key(id)] {
match t.kv_get(&k) {
Ok(Some(v)) => out.push(op(KV_PUT, k.as_bytes(), None, Some(&v))),
_ => out.push(op(KV_DEL, k.as_bytes(), None, None)),
}
}
let done = list_key(id, "done");
out.push(op(SET_CLEAR, done.as_bytes(), None, None));
for m in t.set_members(&done).unwrap_or_default() {
out.push(op(SET_INSERT, done.as_bytes(), None, Some(&m)));
}
}
for (data, exp) in crate::attachments::stored(&t) {
if data.1.len() <= self.max_op() {
out.push(op(KV_PUT, data.0.as_bytes(), None, Some(&data.1)));
out.push(op(KV_PUT, exp.0.as_bytes(), None, Some(&exp.1)));
}
}
out
}
}
pub fn apply_op(store: &Store, raw: &[u8]) -> pagesift_store::Result<()> {
let Some((&tag, rest)) = raw.split_first() else {
return Ok(());
};
let mut r = Reader(rest);
let Some(a) = r.bytes() else {
return Ok(());
};
let t = store.tasks();
match tag {
KV_PUT => {
if let Some(v) = r.bytes() {
t.kv_put(a, v)?;
}
}
KV_DEL => {
t.kv_delete(a)?;
}
LIST_PUSH => {
if let (Some(i), Some(v)) = (r.u64(), r.bytes())
&& t.list_next(a) <= i
{
t.list_push(a, v)?;
}
}
LIST_CLEAR => {
t.list_clear(a)?;
}
SET_INSERT => {
if let Some(v) = r.bytes() {
t.set_insert(a, v)?;
}
}
SET_CLEAR => {
t.set_clear(a)?;
}
LIST_POP => {
t.list_pop_front(a)?;
}
_ => {}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
struct Vec2(Mutex<Vec<Vec<u8>>>);
impl JobLog for Vec2 {
fn append(&self, op: Vec<u8>) -> Result<(), String> {
self.0.lock().unwrap().push(op);
Ok(())
}
}
struct Broken(Mutex<usize>);
impl JobLog for Broken {
fn append(&self, _op: Vec<u8>) -> Result<(), String> {
let mut n = self.0.lock().unwrap();
*n += 1;
if *n > 1 {
Err("journal full".into())
} else {
Ok(())
}
}
}
#[test]
fn a_failed_append_is_reported_to_the_writer() {
let a = tempfile::tempdir().unwrap();
let src = Store::open_dir(a.path()).unwrap();
let log = Arc::new(Broken(Mutex::new(0)));
let m = JobMirror::new(log.clone());
let r = write(&src, Some(&m), |w| {
w.list_push(&list_key("j1", "pages"), b"a")?;
w.list_push(&list_key("j1", "pages"), b"b")?;
w.set_insert(&list_key("j1", "done"), b"a")
});
let e = r.unwrap_err();
assert!(unreplicated(&e));
let e = e.to_string();
assert!(e.contains("journal full"), "{e}");
assert_eq!(m.failures(), 1);
assert_eq!(*log.0.lock().unwrap(), 2);
}
#[test]
fn replicated_writes_rebuild_the_same_job_records() {
let (a, b) = (tempfile::tempdir().unwrap(), tempfile::tempdir().unwrap());
let (src, dst) = (
Store::open_dir(a.path()).unwrap(),
Store::open_dir(b.path()).unwrap(),
);
let log = Arc::new(Vec2(Mutex::default()));
let m = JobMirror::new(log.clone());
write(&src, Some(&m), |w| {
w.kv_put(&meta_key("j1"), b"{\"meta\":1}")?;
for i in 0..3u8 {
w.list_push(&list_key("j1", "pages"), &[i])?;
w.set_insert(&list_key("j1", "done"), &[i])?;
}
w.list_push(&list_key("j1", "frontier"), b"u")
})
.unwrap();
let ops = log.0.lock().unwrap().clone();
for _ in 0..2 {
for o in &ops {
apply_op(&dst, o).unwrap();
}
}
for o in m.dump(&src) {
apply_op(&dst, &o).unwrap();
}
let (s, d) = (src.tasks(), dst.tasks());
assert_eq!(
d.kv_get(meta_key("j1")).unwrap(),
s.kv_get(meta_key("j1")).unwrap()
);
for &what in LISTS {
let l = list_key("j1", what);
assert_eq!(
d.list_range(&l, 0, usize::MAX).unwrap(),
s.list_range(&l, 0, usize::MAX).unwrap(),
"{what}"
);
}
assert_eq!(d.set_len(list_key("j1", "done")), 3);
}
}