use std::{
collections::{BTreeMap, BTreeSet, VecDeque},
fmt,
path::{Path, PathBuf},
sync::{Arc, Condvar, Mutex, MutexGuard, PoisonError},
thread::JoinHandle,
time::{Duration, Instant},
};
use reqwest::{header, StatusCode};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use tokio::runtime::Handle;
use crate::{
fs::{sidecar, write_atomic},
proto::{Refusal, ReplicaInfo, ReplicaList},
wal::{self, RawDb},
Error, Result,
};
const MAX_BATCH: u32 = 512;
const CHECKPOINT_AFTER: u32 = 1000;
const WAL_CAP_FRAMES: u32 = 50_000;
const HEAD_WAIT: Duration = Duration::from_secs(2);
const SWEEP_EVERY: Duration = Duration::from_secs(5);
const RECONCILE_EVERY: Duration = Duration::from_secs(60);
const RECONCILE_FIRST: Duration = Duration::from_secs(30);
const BACKOFF_MIN: Duration = Duration::from_secs(1);
const BACKOFF_MAX: Duration = Duration::from_secs(30);
const BASE_TIMEOUT: Duration = Duration::from_secs(3600);
const FRAMES_TIMEOUT: Duration = Duration::from_secs(60);
const CONNECT_TIMEOUT: Duration = Duration::from_secs(3);
const OPEN_TIMEOUT: Duration = Duration::from_secs(5);
const STATE_DIR: &str = ".walship/state";
const SNAPSHOT_DIR: &str = ".walship/snapshot";
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Owner {
Here,
Elsewhere(String),
}
type OwnerFn = dyn Fn(&str) -> Owner + Send + Sync;
type FilterFn = dyn Fn(&str) -> bool + Send + Sync;
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
struct Persisted {
generation: u64,
acked: u32,
needs_base: bool,
#[serde(default)]
folded: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
salt: Option<u64>,
}
impl Persisted {
fn fresh() -> Self {
Self {
needs_base: true,
..Self::default()
}
}
}
struct Job {
persisted: Persisted,
sent: u32,
handle: Option<RawDb>,
consecutive_failures: u32,
consecutive_mismatches: u32,
in_failure: bool,
last_error: Option<String>,
last_ok: Option<Instant>,
head: u32,
away: Option<String>,
}
impl Job {
fn new(persisted: Persisted) -> Self {
Self {
persisted,
sent: 0,
handle: None,
consecutive_failures: 0,
consecutive_mismatches: 0,
in_failure: false,
last_error: None,
last_ok: None,
head: 0,
away: None,
}
}
fn restart(&mut self) {
self.persisted.needs_base = true;
self.sent = 0;
}
fn next_generation(&mut self) {
let p = &mut self.persisted;
p.generation += 1;
p.acked = 0;
p.folded = 0;
p.salt = None;
self.head = 0;
self.restart();
}
fn lost_by_peer(&self, listed: Instant) -> bool {
self.away.is_none()
&& !self.persisted.needs_base
&& self.last_ok.is_none_or(|ok| ok < listed)
}
}
struct Track {
epoch: u64,
busy: bool,
queued: bool,
rerun: bool,
missing_since: Option<Instant>,
retry_at: Option<Instant>,
last_step: Option<Instant>,
job: Option<Box<Job>>,
last_seen: Option<DbStatus>,
}
#[derive(Default)]
struct PeerView {
at: Option<Instant>,
replicas: usize,
missing: usize,
orphans: usize,
error: Option<String>,
}
#[derive(Default)]
struct Inner {
tracks: BTreeMap<String, Track>,
queue: VecDeque<String>,
stopping: bool,
peer: PeerView,
}
enum Outcome {
Idle,
More,
Retry(Duration),
Gone,
}
pub struct Builder {
data_dir: PathBuf,
peer: String,
token: Option<String>,
workers: usize,
owner: Option<Arc<OwnerFn>>,
filter: Option<Arc<FilterFn>>,
}
impl fmt::Debug for Builder {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Builder")
.field("data_dir", &self.data_dir)
.field("peer", &self.peer)
.field("workers", &self.workers)
.finish_non_exhaustive()
}
}
impl Builder {
#[must_use]
pub fn token(mut self, token: impl Into<String>) -> Self {
self.token = Some(token.into());
self
}
#[must_use]
pub fn workers(mut self, n: usize) -> Self {
self.workers = n.max(1);
self
}
#[must_use]
pub fn owner<F>(mut self, f: F) -> Self
where
F: Fn(&str) -> Owner + Send + Sync + 'static,
{
self.owner = Some(Arc::new(f));
self
}
#[must_use]
pub fn filter<F>(mut self, f: F) -> Self
where
F: Fn(&str) -> bool + Send + Sync + 'static,
{
self.filter = Some(Arc::new(f));
self
}
pub fn start(self, rt: Handle) -> Result<Shipper> {
let peer = self.peer.trim().trim_end_matches('/').to_owned();
if !(peer.starts_with("http://") || peer.starts_with("https://")) {
return Err(Error::InvalidPeerUrl(peer));
}
let state_dir = self.data_dir.join(STATE_DIR);
std::fs::create_dir_all(&state_dir)?;
let snapshot_dir = self.data_dir.join(SNAPSHOT_DIR);
let _ = std::fs::remove_dir_all(&snapshot_dir);
std::fs::create_dir_all(&snapshot_dir)?;
let mut headers = header::HeaderMap::new();
if let Some(token) = &self.token {
let mut value = header::HeaderValue::try_from(format!("Bearer {token}"))
.map_err(|_| Error::InvalidToken)?;
value.set_sensitive(true);
headers.insert(header::AUTHORIZATION, value);
}
let client = reqwest::Client::builder()
.connect_timeout(CONNECT_TIMEOUT)
.default_headers(headers)
.build()?;
let shared = Arc::new(Shared {
data_dir: self.data_dir,
state_dir,
snapshot_dir,
peer,
client,
rt,
owner: self.owner,
filter: self.filter,
started: Instant::now(),
inner: Mutex::default(),
work: Condvar::new(),
tick: Condvar::new(),
});
let mut threads = Vec::with_capacity(self.workers + 1);
for i in 0..self.workers {
let s = Arc::clone(&shared);
threads.push(
std::thread::Builder::new()
.name(format!("walship-{i}"))
.spawn(move || s.worker())?,
);
}
let s = Arc::clone(&shared);
threads.push(
std::thread::Builder::new()
.name("walship-sweep".into())
.spawn(move || s.sweeper())?,
);
Ok(Shipper { shared, threads })
}
}
pub struct Shipper {
shared: Arc<Shared>,
threads: Vec<JoinHandle<()>>,
}
impl fmt::Debug for Shipper {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Shipper")
.field("data_dir", &self.shared.data_dir)
.field("peer", &self.shared.peer)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct DbStatus {
pub name: String,
pub generation: u64,
pub acked: u32,
pub head: u32,
pub pending: u32,
pub needs_base: bool,
pub failures: u32,
pub stepping: bool,
pub owner: Option<String>,
pub last_error: Option<String>,
pub ok_secs_ago: u64,
pub retry_in_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct PeerStatus {
pub seen_secs_ago: u64,
pub replicas: usize,
pub missing: usize,
pub orphans: usize,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ShipperStatus {
pub peer_url: String,
pub databases: Vec<DbStatus>,
pub peer: Option<PeerStatus>,
pub peer_error: Option<String>,
}
impl Shipper {
pub fn builder(data_dir: impl Into<PathBuf>, peer_url: impl Into<String>) -> Builder {
Builder {
data_dir: data_dir.into(),
peer: peer_url.into(),
token: None,
workers: 2,
owner: None,
filter: None,
}
}
#[must_use]
pub fn peer_url(&self) -> &str {
&self.shared.peer
}
pub fn track(&self, name: &str) {
self.shared.track(name);
}
pub fn wake(&self, name: &str) {
self.shared.wake(name);
}
pub fn reset(&self, name: &str) {
self.shared.reset(name);
}
pub fn forget(&self, name: &str) {
self.shared.forget(name);
}
pub fn reconcile(&self) {
self.shared.reconcile();
}
#[must_use]
pub fn status(&self) -> ShipperStatus {
self.shared.status()
}
pub fn shutdown(mut self) {
self.shared.stop();
for t in self.threads.drain(..) {
let _ = t.join();
}
}
}
impl Drop for Shipper {
fn drop(&mut self) {
self.shared.stop();
}
}
pub fn mark_needs_base(data_dir: &Path, name: &str) -> Result<()> {
let path = data_dir.join(STATE_DIR).join(format!("{name}.json"));
let Ok(bytes) = std::fs::read(&path) else {
return Ok(());
};
let mut p: Persisted = serde_json::from_slice(&bytes)?;
if !p.needs_base {
p.needs_base = true;
write_atomic(&path, &serde_json::to_vec(&p)?)?;
}
Ok(())
}
struct Snapshot {
path: PathBuf,
}
impl Snapshot {
fn take(db: &RawDb, dir: &Path, name: &str) -> Result<Self> {
std::fs::create_dir_all(dir)?;
let snapshot = Self {
path: dir.join(format!("{name}.snapshot")),
};
db.copy_main_file(&snapshot.path)?;
Ok(snapshot)
}
}
impl Drop for Snapshot {
fn drop(&mut self) {
let _ = std::fs::remove_file(&self.path);
}
}
struct Shared {
data_dir: PathBuf,
state_dir: PathBuf,
snapshot_dir: PathBuf,
peer: String,
client: reqwest::Client,
rt: Handle,
owner: Option<Arc<OwnerFn>>,
filter: Option<Arc<FilterFn>>,
started: Instant,
inner: Mutex<Inner>,
work: Condvar,
tick: Condvar,
}
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(PoisonError::into_inner)
}
async fn send<T: DeserializeOwned>(
req: reqwest::RequestBuilder,
) -> Result<(StatusCode, Option<T>)> {
let resp = req.send().await?;
let status = resp.status();
let body = resp.bytes().await?;
Ok((status, serde_json::from_slice(&body).ok()))
}
impl Shared {
fn ships(&self, name: &str) -> bool {
self.filter.as_ref().is_none_or(|f| f(name))
}
fn owner_of(&self, name: &str) -> Owner {
self.owner.as_ref().map_or(Owner::Here, |f| f(name))
}
fn state_path(&self, name: &str) -> PathBuf {
self.state_dir.join(format!("{name}.json"))
}
fn db_path(&self, name: &str) -> PathBuf {
self.data_dir.join(format!("{name}.db"))
}
fn replica_url(&self, name: &str) -> String {
format!("{}/replica/{name}", self.peer)
}
fn load(&self, name: &str) -> Persisted {
std::fs::read(self.state_path(name))
.ok()
.and_then(|b| serde_json::from_slice(&b).ok())
.unwrap_or_else(Persisted::fresh)
}
fn persist(&self, name: &str, p: &Persisted) -> Result<()> {
write_atomic(&self.state_path(name), &serde_json::to_vec(p)?)?;
Ok(())
}
fn stop(&self) {
lock(&self.inner).stopping = true;
self.work.notify_all();
self.tick.notify_all();
}
fn enqueue(&self, inner: &mut Inner, name: &str) {
let Some(t) = inner.tracks.get_mut(name) else {
return;
};
if !t.queued && !t.busy {
t.queued = true;
inner.queue.push_back(name.to_owned());
self.work.notify_one();
}
}
fn track(&self, name: &str) {
if !self.ships(name) {
return;
}
let mut inner = lock(&self.inner);
if !inner.tracks.contains_key(name) {
let job = Job::new(self.load(name));
let last_seen = Some(self.describe(name, &job, Instant::now()));
inner.tracks.insert(
name.to_owned(),
Track {
epoch: 0,
busy: false,
queued: false,
rerun: false,
missing_since: None,
retry_at: None,
last_step: None,
last_seen,
job: Some(Box::new(job)),
},
);
}
self.enqueue(&mut inner, name);
}
fn wake(&self, name: &str) {
if !self.ships(name) {
return;
}
let mut inner = lock(&self.inner);
let Some(t) = inner.tracks.get_mut(name) else {
return;
};
if t.retry_at.is_some_and(|at| at > Instant::now()) {
return;
}
if t.busy {
t.rerun = true;
return;
}
self.enqueue(&mut inner, name);
}
fn reset(&self, name: &str) {
if !self.ships(name) {
return;
}
self.track(name);
let mut inner = lock(&self.inner);
let Some(t) = inner.tracks.get_mut(name) else {
return;
};
t.epoch += 1;
t.retry_at = None;
match t.job.as_mut() {
Some(job) => {
job.handle = None;
job.restart();
let _ = self.persist(name, &job.persisted);
}
None => t.rerun = true,
}
self.enqueue(&mut inner, name);
}
fn forget(&self, name: &str) {
if !self.ships(name) {
return;
}
{
let mut inner = lock(&self.inner);
inner.tracks.remove(name);
inner.queue.retain(|n| n != name);
}
let _ = std::fs::remove_file(self.state_path(name));
self.drop_peer_replica(name);
}
fn drop_peer_replica(&self, name: &str) {
let req = self
.client
.delete(self.replica_url(name))
.timeout(FRAMES_TIMEOUT);
self.rt.spawn(async move {
let _ = req.send().await;
});
}
fn status(&self) -> ShipperStatus {
let inner = lock(&self.inner);
let now = Instant::now();
let databases = inner
.tracks
.iter()
.filter_map(|(name, t)| {
let mut row = if let Some(job) = &t.job {
self.describe(name, job, now)
} else {
let mut row = t.last_seen.clone()?;
row.stepping = true;
row
};
row.retry_in_secs = t
.retry_at
.map(|at| at.saturating_duration_since(now).as_secs());
Some(row)
})
.collect();
let peer = inner.peer.at.map(|at| PeerStatus {
seen_secs_ago: now.duration_since(at).as_secs(),
replicas: inner.peer.replicas,
missing: inner.peer.missing,
orphans: inner.peer.orphans,
});
ShipperStatus {
peer_url: self.peer.clone(),
databases,
peer,
peer_error: inner.peer.error.clone(),
}
}
fn describe(&self, name: &str, job: &Job, now: Instant) -> DbStatus {
let p = &job.persisted;
DbStatus {
name: name.to_owned(),
generation: p.generation,
acked: p.acked,
head: job.head,
pending: job.head.saturating_sub(p.acked),
needs_base: p.needs_base,
failures: job.consecutive_failures,
stepping: false,
owner: job.away.clone(),
last_error: job.last_error.clone(),
ok_secs_ago: now
.duration_since(job.last_ok.unwrap_or(self.started))
.as_secs(),
retry_in_secs: None,
}
}
fn sweeper(&self) {
let mut next_reconcile = Instant::now() + RECONCILE_FIRST;
let mut inner = lock(&self.inner);
loop {
inner = self
.tick
.wait_timeout(inner, Duration::from_secs(1))
.unwrap_or_else(PoisonError::into_inner)
.0;
if inner.stopping {
return;
}
if Instant::now() >= next_reconcile {
next_reconcile = Instant::now() + RECONCILE_EVERY;
drop(inner);
self.reconcile();
inner = lock(&self.inner);
if inner.stopping {
return;
}
}
let now = Instant::now();
let due: Vec<String> = inner
.tracks
.iter()
.filter(|(_, t)| !t.busy && !t.queued)
.filter(|(_, t)| match t.retry_at {
Some(at) => at <= now,
None => t
.last_step
.is_none_or(|l| now.duration_since(l) >= SWEEP_EVERY),
})
.map(|(name, _)| name.clone())
.collect();
for name in due {
if let Some(t) = inner.tracks.get_mut(&name) {
t.retry_at = None;
}
self.enqueue(&mut inner, &name);
}
}
}
fn reconcile(&self) {
let started = Instant::now();
let names = match self.peer_replica_names() {
Ok(names) => names,
Err(e) => {
lock(&self.inner).peer.error = Some(e.to_string());
return;
}
};
let mut inner = lock(&self.inner);
let mut missing = 0;
let mut rebase = Vec::new();
for (name, t) in &mut inner.tracks {
if names.contains(name) {
continue;
}
let Some(job) = t.job.as_mut() else {
t.missing_since = Some(started);
missing += 1;
continue;
};
if job.away.is_some() {
continue;
}
missing += 1;
if job.lost_by_peer(started) {
job.restart();
let _ = self.persist(name, &job.persisted);
rebase.push(name.clone());
}
}
let orphans = names
.iter()
.filter(|n| !inner.tracks.contains_key(*n))
.count();
inner.peer = PeerView {
at: Some(started),
replicas: names.len(),
missing,
orphans,
error: None,
};
if !rebase.is_empty() {
tracing::info!(
peer = %self.peer,
"{} database(s) have no replica on the peer, sending bases: {}",
rebase.len(),
rebase.join(", ")
);
}
for name in &rebase {
self.enqueue(&mut inner, name);
}
}
fn peer_replica_names(&self) -> Result<BTreeSet<String>> {
let req = self
.client
.get(format!("{}/replica", self.peer))
.timeout(FRAMES_TIMEOUT);
let (status, list) = self.rt.block_on(send::<ReplicaList>(req))?;
match list {
Some(list) if status.is_success() => {
Ok(list.replicas.into_iter().map(|r| r.name).collect())
}
_ => Err(Error::Peer(format!("GET /replica answered {status}"))),
}
}
fn worker(&self) {
loop {
let (name, epoch, mut job) = {
let mut inner = lock(&self.inner);
loop {
if inner.stopping {
return;
}
if let Some(name) = inner.queue.pop_front() {
let Some(t) = inner.tracks.get_mut(&name) else {
continue;
};
t.queued = false;
if t.busy {
continue;
}
let Some(job) = t.job.take() else { continue };
t.busy = true;
t.rerun = false;
break (name, t.epoch, job);
}
inner = self
.work
.wait(inner)
.unwrap_or_else(PoisonError::into_inner);
}
};
let mut outcome = self.step(&name, &mut job);
let mut inner = lock(&self.inner);
let now = Instant::now();
let row = self.describe(&name, &job, now);
let Some(t) = inner.tracks.get_mut(&name) else {
continue;
};
t.busy = false;
t.last_step = Some(now);
t.last_seen = Some(row);
if t.epoch != epoch {
job.handle = None;
job.restart();
let _ = self.persist(&name, &job.persisted);
t.retry_at = None;
t.job = Some(job);
self.enqueue(&mut inner, &name);
continue;
}
let rerun = std::mem::take(&mut t.rerun);
if let Some(listed) = t.missing_since.take() {
if !matches!(outcome, Outcome::Gone) && job.lost_by_peer(listed) {
tracing::info!(db = %name, peer = %self.peer, "no replica on the peer, sending a base");
job.restart();
let _ = self.persist(&name, &job.persisted);
t.retry_at = None;
outcome = Outcome::More;
}
}
match outcome {
Outcome::Gone => {
inner.tracks.remove(&name);
}
Outcome::Retry(after) => {
t.retry_at = Some(now + after);
t.job = Some(job);
}
Outcome::More => {
t.job = Some(job);
self.enqueue(&mut inner, &name);
}
Outcome::Idle => {
t.job = Some(job);
if rerun {
self.enqueue(&mut inner, &name);
}
}
}
}
}
fn step(&self, name: &str, job: &mut Job) -> Outcome {
let path = self.db_path(name);
if !path.exists() {
return Outcome::Gone;
}
self.step_inner(name, &path, job)
.unwrap_or_else(|e| self.failed(name, job, &e))
}
fn handle<'a>(path: &Path, job: &'a mut Job) -> Result<&'a RawDb> {
let db = if let Some(db) = job.handle.take() {
db
} else {
let db = RawDb::open(path, OPEN_TIMEOUT)?;
db.disable_checkpoint()?;
db
};
Ok(job.handle.insert(db))
}
fn step_inner(&self, name: &str, path: &Path, job: &mut Job) -> Result<Outcome> {
if self.ownership_moved_away(name, job) {
return Ok(Outcome::Idle);
}
if job.persisted.needs_base {
return self.rebase(name, path, job);
}
let Some(head) = Self::handle(path, job)?.head(HEAD_WAIT)? else {
return Ok(Outcome::Retry(Duration::from_millis(200)));
};
job.head = head;
let salt = wal::wal_salt(&sidecar(path, "-wal"))?.map(u64::from_be_bytes);
let p = &mut job.persisted;
let salt_changed = matches!((p.salt, salt), (Some(stored), Some(cur)) if stored != cur);
if salt_changed || head < p.acked {
if salt_changed && p.folded == p.acked && p.acked > 0 {
p.generation += 1;
p.acked = 0;
p.folded = 0;
p.salt = salt;
job.sent = 0;
self.persist(name, p)?;
} else {
p.needs_base = true;
self.persist(name, p)?;
return Ok(Outcome::More);
}
}
if p.salt.is_none() && salt.is_some() && head > 0 {
p.salt = salt;
self.persist(name, p)?;
}
if head >= WAL_CAP_FRAMES {
p.needs_base = true;
self.persist(name, p)?;
return Ok(Outcome::More);
}
if head == p.acked {
job.sent = head;
if head >= CHECKPOINT_AFTER && p.folded != head {
return self.checkpoint(name, job, head);
}
return Ok(Outcome::Idle);
}
if job.sent >= head {
job.sent = p.acked;
}
self.ship(name, path, job, head)
}
fn ownership_moved_away(&self, name: &str, job: &mut Job) -> bool {
match (self.owner_of(name), job.away.take()) {
(Owner::Elsewhere(owner), Some(was)) => {
if was != owner {
tracing::info!(db = name, "owned by {owner} (was {was}), not shipping");
}
job.away = Some(owner);
true
}
(Owner::Elsewhere(owner), None) => {
tracing::info!(
db = name,
peer = %self.peer,
"owned by {owner}, deleting the replica on the peer"
);
self.drop_peer_replica(name);
let p = &mut job.persisted;
p.acked = 0;
p.folded = 0;
p.salt = None;
job.restart();
job.head = 0;
job.handle = None;
job.last_error = None;
let _ = self.persist(name, &job.persisted);
job.away = Some(owner);
true
}
(Owner::Here, Some(was)) => {
tracing::info!(db = name, "owned here again (was {was}), sending a base");
job.restart();
false
}
(Owner::Here, None) => false,
}
}
fn rebase(&self, name: &str, path: &Path, job: &mut Job) -> Result<Outcome> {
let db = Self::handle(path, job)?;
let ck = db.checkpoint_truncate()?;
if !ck.complete() {
job.last_error = Some(format!(
"checkpoint before base incomplete ({}/{} frames, busy={})",
ck.checkpointed, ck.log, ck.busy
));
return Ok(Outcome::Retry(BACKOFF_MIN));
}
let snapshot = Snapshot::take(db, &self.snapshot_dir, name)?;
job.next_generation();
self.persist(name, &job.persisted)?;
let generation = job.persisted.generation;
let url = format!("{}/base?generation={generation}", self.replica_url(name));
let (status, info) = self.rt.block_on(async {
let file = tokio::fs::File::open(&snapshot.path).await?;
let len = file.metadata().await?.len();
let req = self
.client
.post(&url)
.header(header::CONTENT_LENGTH, len)
.header(header::CONTENT_TYPE, "application/octet-stream")
.timeout(BASE_TIMEOUT)
.body(reqwest::Body::wrap_stream(
tokio_util::io::ReaderStream::new(file),
));
send::<ReplicaInfo>(req).await
})?;
let installed = match info {
Some(info) if status.is_success() => info.generation,
_ => return Err(Error::Peer(format!("base refused with {status}"))),
};
if installed != generation {
return Err(Error::Peer(format!(
"base installed as generation {installed}, sent {generation}"
)));
}
job.persisted.needs_base = false;
self.persist(name, &job.persisted)?;
self.recovered(name, job);
tracing::info!(db = name, "rebased at generation {generation}");
Ok(Outcome::More)
}
fn checkpoint(&self, name: &str, job: &mut Job, head: u32) -> Result<Outcome> {
let db = job.handle.as_ref().expect("step_inner opened the handle");
let ck = db.checkpoint_passive()?;
let p = &mut job.persisted;
if ck.checkpointed > head {
p.needs_base = true;
self.persist(name, p)?;
return Ok(Outcome::More);
}
if ck.log == head && ck.checkpointed == head {
p.folded = head;
self.persist(name, p)?;
return Ok(Outcome::Idle);
}
if ck.log > head {
return Ok(Outcome::More);
}
job.last_error = Some(format!(
"checkpoint incomplete ({}/{} frames, busy={})",
ck.checkpointed, ck.log, ck.busy
));
Ok(Outcome::Idle)
}
fn ship(&self, name: &str, path: &Path, job: &mut Job, head: u32) -> Result<Outcome> {
let from = job.sent.max(job.persisted.acked) + 1;
let to = head.min(from + MAX_BATCH - 1);
let db = Self::handle(path, job)?;
let page_size = db.page_size()?;
let frame_len = wal::FRAME_HEADER + page_size as usize;
let mut body = Vec::with_capacity((to - from + 1) as usize * frame_len);
for no in from..=to {
body.extend(db.get_frame(no, page_size)?);
}
let generation = job.persisted.generation;
let req = self
.client
.post(format!(
"{}/frames?generation={generation}&first={from}&page_size={page_size}",
self.replica_url(name)
))
.header(header::CONTENT_TYPE, "application/octet-stream")
.timeout(FRAMES_TIMEOUT)
.body(body);
let (status, reply) = self.rt.block_on(send::<serde_json::Value>(req))?;
let reply = reply.unwrap_or_default();
match status {
StatusCode::OK => {
let info: ReplicaInfo = serde_json::from_value(reply)?;
job.sent = to;
job.consecutive_mismatches = 0;
if info.frames != job.persisted.acked {
job.persisted.acked = info.frames;
self.persist(name, &job.persisted)?;
}
self.recovered(name, job);
Ok(if to < head {
Outcome::More
} else {
Outcome::Idle
})
}
StatusCode::CONFLICT => {
let refusal: Refusal = serde_json::from_value(reply).unwrap_or_else(|_| Refusal {
needs_base: true,
..Refusal::new("conflict")
});
job.consecutive_mismatches += 1;
let resume_at = refusal.frames.filter(|&f| {
!refusal.needs_base
&& job.consecutive_mismatches == 1
&& refusal.generation == Some(generation)
&& f <= head
});
if let Some(frames) = resume_at {
job.persisted.acked = frames;
job.sent = frames;
self.persist(name, &job.persisted)?;
return Ok(Outcome::More);
}
job.consecutive_mismatches = 0;
job.last_error = Some(format!(
"peer refused frames {from}..={to}: {}",
refusal.error
));
job.persisted.needs_base = true;
self.persist(name, &job.persisted)?;
Ok(Outcome::More)
}
StatusCode::SERVICE_UNAVAILABLE => Ok(Outcome::Retry(BACKOFF_MIN)),
_ => {
let error = reply.get("error").and_then(|e| e.as_str()).unwrap_or("");
Err(Error::Peer(format!(
"{status} for frames {from}..={to}: {error}"
)))
}
}
}
fn failed(&self, name: &str, job: &mut Job, e: &Error) -> Outcome {
job.consecutive_failures = job.consecutive_failures.saturating_add(1);
job.last_error = Some(e.to_string());
if !job.in_failure {
job.in_failure = true;
tracing::warn!(db = name, peer = %self.peer, "can't ship: {e}");
}
if job.consecutive_failures >= 3 && !job.persisted.needs_base {
self.cap_wal(name, job);
}
let exp = job.consecutive_failures.min(6) - 1;
Outcome::Retry((BACKOFF_MIN * 2u32.pow(exp)).min(BACKOFF_MAX))
}
fn cap_wal(&self, name: &str, job: &mut Job) {
let Some(db) = job.handle.as_ref() else {
return;
};
let Ok(Some(head)) = db.head(Duration::from_millis(200)) else {
return;
};
if head.saturating_sub(job.persisted.acked) <= WAL_CAP_FRAMES {
return;
}
if !db.checkpoint_truncate().is_ok_and(|ck| ck.complete()) {
return;
}
job.next_generation();
let _ = self.persist(name, &job.persisted);
tracing::warn!(
db = name,
"WAL passed {WAL_CAP_FRAMES} frames, folded it; the peer gets a base"
);
}
fn recovered(&self, name: &str, job: &mut Job) {
job.last_ok = Some(Instant::now());
job.consecutive_failures = 0;
job.last_error = None;
if std::mem::take(&mut job.in_failure) {
tracing::info!(db = name, peer = %self.peer, "shipping again");
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn mark_needs_base_flags_the_state_file_and_skips_a_missing_one() {
let dir = std::env::temp_dir().join(format!("walship-mark-{}", std::process::id()));
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(dir.join(STATE_DIR)).unwrap();
mark_needs_base(&dir, "x").unwrap();
assert!(!dir.join(STATE_DIR).join("x.json").exists());
let path = dir.join(STATE_DIR).join("x.json");
let p = Persisted {
generation: 4,
acked: 17,
salt: Some(0xdead_beef),
..Persisted::default()
};
std::fs::write(&path, serde_json::to_vec(&p).unwrap()).unwrap();
mark_needs_base(&dir, "x").unwrap();
let back: Persisted = serde_json::from_slice(&std::fs::read(&path).unwrap()).unwrap();
assert_eq!(
back,
Persisted {
needs_base: true,
..p
}
);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn a_peer_url_without_a_scheme_is_refused() {
let rt = tokio::runtime::Runtime::new().unwrap();
let err = Shipper::builder(std::env::temp_dir(), "standby.internal:7350")
.start(rt.handle().clone())
.unwrap_err();
assert!(matches!(err, Error::InvalidPeerUrl(_)), "{err}");
}
}