mod fsio;
mod manifest;
pub use crate::ObjectStore;
pub use manifest::{Lease, Segment};
use crate::{fenced, Bytes};
use graph_r::format::{base, wal};
use graph_r::fsio::now_ms;
use graph_r::hash::xxh3_64;
use graph_r::store::{BASE_FILE, WAL_FILE};
use graph_r::Store;
use manifest::Manifest;
use std::future::Future;
use std::io::{Read, Seek, SeekFrom};
use std::path::Path;
use std::sync::{Mutex, MutexGuard, PoisonError};
pub const COALESCE_SEGMENTS: usize = 32;
#[doc(hidden)]
pub use manifest::fuzz as fuzz_manifest;
pub type Result<T, E = Error> = std::result::Result<T, E>;
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
#[error(transparent)]
Graph(#[from] graph_r::Error),
#[error("object store: {0}")]
Store(crate::Error),
#[error("fenced: {reason}")]
Fenced {
reason: &'static str,
},
#[error("lease held by {holder} until {expires_ms}")]
Leased {
holder: String,
expires_ms: u64,
},
#[error("corrupt remote store: {message}")]
Corrupt {
message: String,
},
#[error("i/o: {0}")]
Io(#[from] std::io::Error),
}
fn corrupt(message: impl Into<String>) -> Error {
Error::Corrupt { message: message.into() }
}
impl From<crate::Error> for Error {
fn from(e: crate::Error) -> Self {
match e {
crate::Error::Fenced { reason } => Self::Fenced { reason },
crate::Error::Corrupt { message } => Self::Corrupt { message },
other => Self::Store(other),
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Pulled {
pub generation: u64,
pub bytes: u64,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Published {
pub generation: u64,
pub wal_len: u64,
pub uploaded: u64,
pub new_base: bool,
}
enum Change {
Tail(Vec<Bytes>),
Full(Bytes, Vec<Bytes>),
}
#[derive(Clone, Debug)]
struct Remote {
head_gen: u64,
manifest: Manifest,
local_verified: bool,
}
#[derive(Debug)]
struct Holder {
name: String,
ttl_ms: u64,
}
#[derive(Debug, Default)]
struct State {
remote: Option<Remote>,
holder: Option<Holder>,
image_wal: Vec<u8>,
}
enum Cut {
Nothing { generation: u64, wal_len: u64 },
Tail { generation: u64, start: u64, frames: Vec<u8> },
Coalesce { generation: u64, wal: Vec<u8> },
Full { generation: u64, base: Vec<u8>, wal: Vec<u8> },
}
#[derive(Debug)]
pub struct Mirror<'o, O: ObjectStore> {
objects: &'o O,
prefix: String,
state: Mutex<State>,
}
impl<'o, O: ObjectStore> Mirror<'o, O> {
pub fn new(objects: &'o O, prefix: &str) -> Result<Self> {
let prefix = fenced::prefix(prefix)?;
Ok(Self { objects, prefix, state: Mutex::new(State::default()) })
}
#[must_use]
pub fn prefix(&self) -> &str {
&self.prefix
}
#[must_use]
pub fn lease(&self) -> Option<Lease> {
self.lock().remote.as_ref().and_then(|r| r.manifest.lease.clone())
}
fn lock(&self) -> MutexGuard<'_, State> {
self.state.lock().unwrap_or_else(PoisonError::into_inner)
}
fn head_key(&self) -> String {
format!("{}HEAD", self.prefix)
}
fn base_key(&self, generation: u64) -> String {
format!("{}base/{generation:016x}.base", self.prefix)
}
fn gen_dir(&self, generation: u64) -> String {
format!("{}wal/{generation:016x}/", self.prefix)
}
fn seg_key(&self, generation: u64, s: &Segment) -> String {
format!("{}{:016x}-{:016x}.seg", self.gen_dir(generation), s.start, s.end())
}
async fn fetch_remote(&self) -> Result<Option<Remote>> {
let Some(obj) = self.objects.get(&self.head_key()).await? else {
self.lock().remote = None;
return Ok(None);
};
let manifest = Manifest::decode(&obj.bytes)?;
let mut st = self.lock();
let local_verified =
st.remote.as_ref().is_some_and(|r| r.local_verified && r.head_gen == obj.generation);
let remote = Remote { head_gen: obj.generation, manifest, local_verified };
st.remote = Some(remote.clone());
Ok(Some(remote))
}
fn remember(&self, remote: Option<Remote>) {
self.lock().remote = remote;
}
async fn cas_head(&self, prev: Option<&Remote>, next: &Manifest) -> Result<u64> {
let body = Bytes::from(next.encode());
match fenced::cas(self.objects, &self.head_key(), prev.map(|r| r.head_gen), body).await {
Ok(g) => Ok(g),
Err(e) => {
self.remember(None);
Err(e.into())
}
}
}
async fn fetch_all(&self, m: &Manifest) -> Result<(Bytes, Vec<Bytes>)> {
let base_key = self.base_key(m.generation);
let base = async {
Ok(fenced::fetch_verified(self.objects, &base_key, m.base_len, m.base_sum).await?)
};
let segs = self.fetch_segments(m, m.segments.iter());
let (base, segs) = futures_util::future::try_join(base, segs).await?;
let generation = base::generation_of(&base).map_err(|e| corrupt(format!("base: {e}")))?;
if generation != m.generation {
return Err(corrupt("base generation differs from the manifest"));
}
Ok((base, segs))
}
async fn fetch_segments<'s>(
&'s self,
m: &'s Manifest,
wanted: impl Iterator<Item = &'s Segment>,
) -> Result<Vec<Bytes>> {
let keys: Vec<(String, Segment)> =
wanted.map(|s| (self.seg_key(m.generation, s), *s)).collect();
futures_util::future::try_join_all(keys.iter().map(|(key, s)| async move {
let b = fenced::fetch_verified(self.objects, key, s.len, s.sum).await?;
check_framing(&b, s.start, m.generation)?;
Ok::<_, Error>(b)
}))
.await
}
pub fn load_image(&self) -> impl Future<Output = Result<Option<Store>>> + '_ {
async move {
let Some(remote) = self.fetch_remote().await? else { return Ok(None) };
if remote.manifest.generation == 0 {
return Ok(None);
}
let (base, segs) = self.fetch_all(&remote.manifest).await?;
let wal = concat(&segs, remote.manifest.wal_len());
let store = Store::open_image(base, &wal)?;
let mut st = self.lock();
st.image_wal = wal;
st.remote = Some(Remote { local_verified: true, ..remote });
Ok(Some(store))
}
}
pub fn refresh_image<'a>(
&'a self,
store: &'a Store,
) -> impl Future<Output = Result<bool>> + 'a {
async move {
let cached = self.lock().remote.clone();
let Some((remote, change)) = self.changes(cached.as_ref(), false, |_| true).await?
else {
return Ok(false);
};
let changed = match change {
Change::Tail(tail) => {
let mut st = self.lock();
let have = cached.as_ref().map_or(0, |c| c.manifest.wal_len());
if st.image_wal.len() as u64 != have {
return Err(corrupt("replica log image out of step with the manifest"));
}
for b in &tail {
st.image_wal.extend_from_slice(b);
}
store.reload_image(None, &st.image_wal)?
}
Change::Full(base, segs) => {
let wal = concat(&segs, remote.manifest.wal_len());
let changed = store.reload_image(Some(base), &wal)?;
self.lock().image_wal = wal;
changed
}
};
self.remember(Some(Remote { local_verified: true, ..remote }));
Ok(changed)
}
}
async fn changes(
&self,
cached: Option<&Remote>,
local: bool,
local_ok: impl Fn(&Remote) -> bool,
) -> Result<Option<(Remote, Change)>> {
let meta = self.objects.head(&self.head_key()).await?;
match (&meta, cached) {
(None, _) => return Ok(None),
(Some(m), Some(c)) if m.generation == c.head_gen && (!local || c.local_verified) => {
return Ok(None)
}
_ => {}
}
let Some(remote) = self.fetch_remote().await? else { return Ok(None) };
let m = &remote.manifest;
if m.generation == 0 {
return Ok(None);
}
let prefix = cached.filter(|c| {
c.manifest.generation == m.generation
&& c.manifest.base_sum == m.base_sum
&& m.wal_len() >= c.manifest.wal_len()
&& local_ok(c)
});
let change = if let Some(c) = prefix {
let have = c.manifest.wal_len();
let wanted = || m.segments.iter().filter(|s| s.end() > have);
let segs = self.fetch_segments(m, wanted()).await?;
Change::Tail(
wanted()
.zip(segs)
.map(|(s, b)| {
let skip =
usize::try_from(have.saturating_sub(s.start)).unwrap_or(usize::MAX);
b.slice(skip.min(b.len())..)
})
.collect(),
)
} else {
let (base, segs) = self.fetch_all(m).await?;
Change::Full(base, segs)
};
Ok(Some((remote, change)))
}
pub fn pull_into<'a>(
&'a self,
dir: &'a Path,
) -> impl Future<Output = Result<Option<Pulled>>> + 'a {
async move {
let Some(remote) = self.fetch_remote().await? else { return Ok(None) };
if remote.manifest.generation == 0 {
return Ok(None);
}
let _lock = fsio::lock_dir(dir)?;
let (base, segs) = self.fetch_all(&remote.manifest).await?;
write_store(dir, &base, &segs)?;
let bytes = base.len() as u64 + remote.manifest.wal_len();
let generation = remote.manifest.generation;
self.remember(Some(Remote { local_verified: true, ..remote }));
Ok(Some(Pulled { generation, bytes }))
}
}
pub fn refresh_into<'a>(&'a self, store: &'a Store) -> impl Future<Output = Result<bool>> + 'a {
async move {
let cached = self.lock().remote.clone();
let dir = store.dir();
let wal_path = dir.join(WAL_FILE);
let _lock = fsio::lock_dir(dir)?;
let local_len = std::fs::metadata(&wal_path).map(|md| md.len()).ok();
let local_ok = |c: &Remote| c.local_verified && local_len == Some(c.manifest.wal_len());
let Some((remote, change)) = self.changes(cached.as_ref(), true, local_ok).await?
else {
return Ok(false);
};
match change {
Change::Tail(tail) => {
let parts: Vec<&[u8]> = tail.iter().map(AsRef::as_ref).collect();
fsio::append(&wal_path, &parts)?;
}
Change::Full(base, segs) => write_store(dir, &base, &segs)?,
}
self.remember(Some(Remote { local_verified: true, ..remote }));
store.reload_if_stale()?;
Ok(true)
}
}
pub fn publish<'a>(&'a self, store: &'a Store) -> impl Future<Output = Result<Published>> + 'a {
async move {
let cached = self.lock().remote.clone();
let remote = match cached {
Some(r) => Some(r),
None => self.fetch_remote().await?,
};
let lease = self.lease_for_write(remote.as_ref())?;
let cut = local_cut(store, remote.as_ref())?;
let prev_segments =
remote.as_ref().map(|r| r.manifest.segments.clone()).unwrap_or_default();
let prev_gen = remote.as_ref().map_or(0, |r| r.manifest.generation);
let mut next = remote.as_ref().map(|r| r.manifest.clone()).unwrap_or_default();
next.lease = lease;
let (uploaded, new_base, stale): (u64, bool, Vec<String>) = match cut {
Cut::Nothing { generation, wal_len } => {
if let Some(r) = remote {
self.remember(Some(Remote { local_verified: true, ..r }));
}
return Ok(Published { generation, wal_len, uploaded: 0, new_base: false });
}
Cut::Tail { generation, start, frames } => {
let s = Segment { start, len: frames.len() as u64, sum: xxh3_64(&frames) };
fenced::put_immutable(
self.objects,
&self.seg_key(generation, &s),
Bytes::from(frames),
)
.await?;
next.segments.push(s);
(s.len, false, Vec::new())
}
Cut::Coalesce { generation, wal } => {
let s = Segment { start: 0, len: wal.len() as u64, sum: xxh3_64(&wal) };
fenced::put_immutable(
self.objects,
&self.seg_key(generation, &s),
Bytes::from(wal),
)
.await?;
let stale = prev_segments.iter().map(|p| self.seg_key(generation, p)).collect();
next.segments = vec![s];
(s.len, false, stale)
}
Cut::Full { generation, base, wal } => {
let base_sum = xxh3_64(&base);
let base_len = base.len() as u64;
fenced::put_immutable(
self.objects,
&self.base_key(generation),
Bytes::from(base),
)
.await?;
let s = Segment { start: 0, len: wal.len() as u64, sum: xxh3_64(&wal) };
fenced::put_immutable(
self.objects,
&self.seg_key(generation, &s),
Bytes::from(wal),
)
.await?;
next.generation = generation;
next.base_len = base_len;
next.base_sum = base_sum;
next.segments = vec![s];
let stale = if prev_gen != 0 && prev_gen != generation {
vec![self.base_key(prev_gen), self.gen_dir(prev_gen)]
} else {
Vec::new()
};
(base_len + s.len, true, stale)
}
};
let head_gen = self.cas_head(remote.as_ref(), &next).await?;
let published = Published {
generation: next.generation,
wal_len: next.wal_len(),
uploaded,
new_base,
};
self.remember(Some(Remote { head_gen, manifest: next, local_verified: true }));
self.cleanup(&stale).await;
Ok(published)
}
}
async fn cleanup(&self, stale: &[String]) {
for key in stale {
if key.ends_with('/') {
if let Ok(list) = self.objects.list(key).await {
for m in list {
let _ = self.objects.delete(&m.key).await;
}
}
} else {
let _ = self.objects.delete(key).await;
}
}
}
fn lease_for_write(&self, remote: Option<&Remote>) -> Result<Option<Lease>> {
let st = self.lock();
let current = remote.and_then(|r| r.manifest.lease.clone());
let now = now_ms();
match (&st.holder, current) {
(Some(h), Some(l)) if l.holder != h.name && l.expires_ms > now => {
Err(Error::Leased { holder: l.holder, expires_ms: l.expires_ms })
}
(Some(h), _) => {
Ok(Some(Lease { holder: h.name.clone(), expires_ms: now.saturating_add(h.ttl_ms) }))
}
(None, Some(l)) if l.expires_ms > now => {
Err(Error::Leased { holder: l.holder, expires_ms: l.expires_ms })
}
(None, other) => Ok(other.filter(|l| l.expires_ms > now)),
}
}
pub fn acquire_lease<'a>(
&'a self,
holder: &'a str,
ttl_ms: u64,
) -> impl Future<Output = Result<Lease>> + 'a {
async move {
if holder.is_empty() || holder.len() > manifest::MAX_HOLDER {
return Err(corrupt("lease holder must be 1..=256 bytes"));
}
let remote = self.fetch_remote().await?;
let now = now_ms();
if let Some(l) = remote.as_ref().and_then(|r| r.manifest.lease.as_ref()) {
if l.holder != holder && l.expires_ms > now {
return Err(Error::Leased {
holder: l.holder.clone(),
expires_ms: l.expires_ms,
});
}
}
let mut next = remote.as_ref().map(|r| r.manifest.clone()).unwrap_or_default();
let lease = Lease { holder: holder.to_owned(), expires_ms: now.saturating_add(ttl_ms) };
next.lease = Some(lease.clone());
let head_gen = self.cas_head(remote.as_ref(), &next).await?;
let verified = remote.as_ref().is_some_and(|r| r.local_verified);
let mut st = self.lock();
st.remote = Some(Remote { head_gen, manifest: next, local_verified: verified });
st.holder = Some(Holder { name: holder.to_owned(), ttl_ms });
Ok(lease)
}
}
pub fn release_lease(&self) -> impl Future<Output = Result<()>> + '_ {
async move {
let Some(name) = self.lock().holder.take().map(|h| h.name) else { return Ok(()) };
let Some(remote) = self.fetch_remote().await? else { return Ok(()) };
if remote.manifest.lease.as_ref().is_none_or(|l| l.holder != name) {
return Ok(());
}
let mut next = remote.manifest.clone();
next.lease = None;
let head_gen = self.cas_head(Some(&remote), &next).await?;
self.remember(Some(Remote {
head_gen,
manifest: next,
local_verified: remote.local_verified,
}));
Ok(())
}
}
}
fn check_framing(b: &[u8], start: u64, generation: u64) -> Result<()> {
let ok = if start == 0 {
wal::parse_header(b).ok() == Some(generation) && wal::committed_len(b) == b.len()
} else {
wal::committed_frames_len(b) == b.len()
};
if ok {
Ok(())
} else {
Err(corrupt("segment is not a run of committed frames"))
}
}
fn concat(segs: &[Bytes], len: u64) -> Vec<u8> {
let mut wal = Vec::with_capacity(usize::try_from(len).unwrap_or(0));
for s in segs {
wal.extend_from_slice(s);
}
wal
}
fn write_store(dir: &Path, base: &Bytes, segs: &[Bytes]) -> Result<()> {
fsio::write_atomic(&dir.join(BASE_FILE), &[base])?;
let parts: Vec<&[u8]> = segs.iter().map(AsRef::as_ref).collect();
fsio::write_atomic(&dir.join(WAL_FILE), &parts)
}
fn local_cut(store: &Store, remote: Option<&Remote>) -> Result<Cut> {
let _writer = store.writer()?;
let dir = store.dir();
let wal_path = dir.join(WAL_FILE);
let mut file = std::fs::File::open(&wal_path)?;
let mut header = [0u8; wal::HEADER_LEN];
file.read_exact(&mut header)?;
let generation = wal::parse_header(&header)?;
let published = remote.filter(|r| r.manifest.generation == generation);
if let Some(r) = remote {
if r.manifest.generation > generation {
return Err(Error::Fenced {
reason: "the published store is newer; pull before writing",
});
}
}
let Some(r) = published else {
let base = std::fs::read(dir.join(BASE_FILE))?;
if base::generation_of(&base)? != generation {
return Err(graph_r::Error::corrupt("base and log generations differ").into());
}
let mut wal = Vec::new();
file.seek(SeekFrom::Start(0))?;
file.read_to_end(&mut wal)?;
wal.truncate(wal::committed_len(&wal));
return Ok(Cut::Full { generation, base, wal });
};
let have = r.manifest.wal_len();
let file_len = file.metadata()?.len();
if file_len < have {
return Err(Error::Fenced {
reason: "the local log is shorter than the published one; pull first",
});
}
let coalesce = r.manifest.segments.len() >= COALESCE_SEGMENTS;
if r.local_verified && !coalesce {
let mut frames = Vec::with_capacity(usize::try_from(file_len - have).unwrap_or(0));
file.seek(SeekFrom::Start(have))?;
file.read_to_end(&mut frames)?;
frames.truncate(wal::committed_frames_len(&frames));
return Ok(if frames.is_empty() {
Cut::Nothing { generation, wal_len: have }
} else {
Cut::Tail { generation, start: have, frames }
});
}
let mut wal = Vec::with_capacity(usize::try_from(file_len).unwrap_or(0));
file.seek(SeekFrom::Start(0))?;
file.read_to_end(&mut wal)?;
if !r.local_verified {
for s in &r.manifest.segments {
let (a, b) = (s.start as usize, s.end() as usize);
if wal.get(a..b).is_none_or(|bytes| xxh3_64(bytes) != s.sum) {
return Err(Error::Fenced {
reason: "the local log diverges from the published one",
});
}
}
}
let committed = wal::HEADER_LEN + wal::committed_frames_len(&wal[wal::HEADER_LEN..]);
wal.truncate(committed);
if committed as u64 == have {
return Ok(Cut::Nothing { generation, wal_len: have });
}
if coalesce {
return Ok(Cut::Coalesce { generation, wal });
}
let frames = wal.split_off(have as usize);
Ok(Cut::Tail { generation, start: have, frames })
}