use crate::envelope::{self, Body, EndSnapshot, Limits, Packer};
use crate::plan::{pending, Delta, Engine, Note, Unsent, RECORD_CHUNK};
use crate::wire::Codec;
use crate::Conn;
use crate::Error;
use std::collections::BTreeMap;
use std::future::Future;
use std::time::Duration;
use uuid::Uuid;
pub trait Transport: Send + Sync + 'static {
fn send(
&self,
replica: Uuid,
message: Vec<u8>,
id: &str,
) -> impl Future<Output = Result<(), Error>> + Send;
}
const PAGE: usize = 20_000;
pub(crate) const BUFFERED: usize = 64 << 20;
const SENT_AT_ONCE: usize = 8 << 20;
#[derive(Clone, Copy, Debug)]
pub(crate) struct Bound {
pub(crate) messages: i64,
pub(crate) bytes: i64,
}
const REFUSALS: i32 = 3;
pub(crate) struct Relay<'a, P> {
pub(crate) engine: &'a Engine,
pub(crate) codecs: &'a [Codec<P>],
pub(crate) limits: Limits,
pub(crate) batch: i64,
pub(crate) timeout: Duration,
pub(crate) backoff: crate::run::Backoff,
pub(crate) bound: Bound,
pub(crate) concurrency: usize,
pub(crate) buffered: usize,
}
#[derive(Debug, Default)]
pub(crate) struct Outcome {
pub(crate) done: usize,
pub(crate) parked: Vec<(i64, String)>,
pub(crate) quarantined: usize,
pub(crate) written: usize,
pub(crate) upserts: usize,
pub(crate) deletes: usize,
pub(crate) collapsed: Vec<Uuid>,
}
#[derive(Debug, Default)]
pub(crate) struct Published {
pub(crate) sent: usize,
pub(crate) failing: Vec<(Uuid, String)>,
}
struct Ready {
replica: Uuid,
body: Body,
}
struct Writer {
epoch: Uuid,
secret: Vec<u8>,
packers: BTreeMap<Uuid, Packer>,
snapshots: BTreeMap<Uuid, String>,
ready: Vec<Ready>,
rows: usize,
ready_bytes: usize,
buffered: usize,
room: usize,
stuck: usize,
last: BTreeMap<Uuid, (i64, i64)>,
began: BTreeMap<Uuid, i64>,
quarantined: bool,
ended: BTreeMap<Uuid, i64>,
written: usize,
upserts: usize,
deletes: usize,
}
impl<P: Sync + 'static> Relay<'_, P> {
pub(crate) async fn record(&self, conn: &mut Conn) -> Result<Outcome, Error> {
let notes = pending(conn, self.batch).await?;
if notes.is_empty() {
return Ok(Outcome::default());
}
if self.engine.order.is_empty() {
let done: Vec<i64> = notes.iter().map(|n| n.id).collect();
crate::plan::mark_done(conn, &done).await?;
return Ok(Outcome {
done: done.len(),
..Outcome::default()
});
}
let (mut delta, notes, parked) = self.plan(conn, notes).await?;
let mut outcome = Outcome {
parked,
..Outcome::default()
};
let (epoch, secret): (Uuid, Vec<u8>) =
sqlx::query_as("SELECT epoch, secret FROM sync.epoch FOR SHARE")
.fetch_one(&mut *conn)
.await?;
let touched = self.engine.touched(conn).await?;
let chains: BTreeMap<Uuid, (i64, i64)> = sqlx::query_as::<_, (Uuid, i64, i64)>(
"SELECT replica_id, last_sequence, last_nonce FROM sync.chains \
WHERE replica_id = ANY($1) AND last_sequence IS NOT NULL",
)
.bind(touched.iter().copied().collect::<Vec<_>>())
.fetch_all(&mut *conn)
.await?
.into_iter()
.map(|(replica, sequence, nonce)| (replica, (sequence, nonce)))
.collect();
let unchained: Vec<Uuid> = touched
.into_iter()
.filter(|r| !delta.snapshots.contains(r) && !chains.contains_key(r))
.collect();
if !unchained.is_empty() {
delta = self.engine.provision(conn, delta, &unchained).await?;
}
let snapshots: Vec<Uuid> = delta.snapshots.iter().copied().collect();
let covered: Vec<i64> = sqlx::query_scalar(
"SELECT id FROM sync.notes WHERE kind = 'resync' AND done_at IS NULL \
AND parked_at IS NULL AND replica_id = ANY($1)",
)
.bind(&snapshots)
.fetch_all(&mut *conn)
.await?;
let written = loop {
sqlx::query("SAVEPOINT sync_write")
.execute(&mut *conn)
.await?;
let written = self.write(conn, &delta, (epoch, &secret), &chains).await?;
let end = match written {
Ok(_) => "RELEASE SAVEPOINT sync_write",
Err(_) => "ROLLBACK TO SAVEPOINT sync_write",
};
sqlx::query(end).execute(&mut *conn).await?;
match written {
Ok(written) => break written,
Err(blocked) => delta = self.engine.replan(conn, delta, blocked).await?,
}
};
(outcome.written, outcome.upserts, outcome.deletes) = written;
let unsent = self.engine.unsent(&delta);
outcome.quarantined = unsent.len();
quarantine(conn, &unsent).await?;
self.engine.settle(conn).await?;
let replicas: Vec<Uuid> = self.engine.touched(conn).await?.into_iter().collect();
let done: Vec<i64> = notes.iter().map(|n| n.id).chain(covered).collect();
crate::plan::mark_done(conn, &done).await?;
outcome.done = notes.len();
let (ids, whys): (Vec<i64>, Vec<String>) = outcome
.parked
.iter()
.map(|(id, why)| (*id, stored(why)))
.unzip();
sqlx::query(
"UPDATE sync.notes n SET parked_at = now(), attempts = attempts + 1, last_error = p.why \
FROM unnest($1::bigint[], $2::text[]) AS p(id, why) WHERE n.id = p.id",
)
.bind(&ids)
.bind(&whys)
.execute(&mut *conn)
.await?;
outcome.collapsed = self.bound(conn, &replicas).await?;
Ok(outcome)
}
async fn plan(
&self,
conn: &mut Conn,
mut notes: Vec<Note>,
) -> Result<(Delta, Vec<Note>, Vec<(i64, String)>), Error> {
let mut parked = Vec::new();
loop {
if notes.is_empty() {
return Ok((self.engine.plan(conn, ¬es).await?, notes, parked));
}
let why = match self.attempt(conn, ¬es).await {
Ok(delta) => return Ok((delta, notes, parked)),
Err(e) if e.is_transient() => return Err(e),
Err(e) => e,
};
let (culprit, why) = self.culprit(conn, ¬es, why).await?;
notes.retain(|n| n.id != culprit);
parked.push((culprit, why));
}
}
async fn culprit(
&self,
conn: &mut Conn,
notes: &[Note],
mut why: Error,
) -> Result<(i64, String), Error> {
let mut suspects = notes.to_vec();
while suspects.len() > 1 {
let right = suspects.split_off(suspects.len() / 2);
match self.attempt(conn, &suspects).await {
Err(e) if e.is_transient() => return Err(e),
Err(e) => {
why = e;
continue;
}
Ok(_) => {}
}
match self.attempt(conn, &right).await {
Err(e) if e.is_transient() => return Err(e),
Err(e) => {
why = e;
suspects = right;
}
Ok(_) => suspects = right.into_iter().last().into_iter().collect(),
}
}
Ok((suspects[0].id, why.to_string()))
}
async fn attempt(&self, conn: &mut Conn, notes: &[Note]) -> Result<Delta, Error> {
sqlx::query("SAVEPOINT sync_plan")
.execute(&mut *conn)
.await?;
let planned = self.engine.plan(conn, notes).await;
let end = if planned.is_ok() {
"RELEASE SAVEPOINT sync_plan"
} else {
"ROLLBACK TO SAVEPOINT sync_plan"
};
sqlx::query(end).execute(&mut *conn).await?;
planned
}
async fn write(
&self,
conn: &mut Conn,
delta: &Delta,
(epoch, secret): (Uuid, &[u8]),
chains: &BTreeMap<Uuid, (i64, i64)>,
) -> Result<Result<(usize, usize, usize), BTreeMap<(Uuid, usize, String), String>>, Error> {
let snapshots: Vec<Uuid> = delta.snapshots.iter().copied().collect();
if !snapshots.is_empty() {
for sql in [
"DELETE FROM sync.held WHERE replica_id = ANY($1)",
"DELETE FROM sync.gone WHERE replica_id = ANY($1)",
] {
sqlx::query(sql)
.bind(&snapshots)
.execute(&mut *conn)
.await?;
}
drop_outbox(conn, &snapshots, i64::MAX, false).await?;
}
let mut dropping = self.engine.dropping(conn).await?;
let ids: BTreeMap<Uuid, String> = snapshots
.iter()
.map(|replica| (*replica, Uuid::now_v7().to_string()))
.collect();
let mut writer = Writer {
epoch,
secret: secret.to_vec(),
packers: self
.engine
.touched(conn)
.await?
.into_iter()
.map(|replica| {
let snapshot = ids.get(&replica).map_or("", String::as_str);
let dropping = dropping.remove(&replica).unwrap_or_default();
(replica, Packer::new(snapshot, self.limits, dropping))
})
.collect(),
snapshots: ids.clone(),
room: self.buffered,
stuck: 0,
ready: Vec::new(),
rows: 0,
ready_bytes: 0,
buffered: 0,
last: chains
.iter()
.filter(|(r, _)| !delta.snapshots.contains(r))
.map(|(r, chain)| (*r, *chain))
.collect(),
began: BTreeMap::new(),
ended: BTreeMap::new(),
quarantined: sqlx::query_scalar("SELECT EXISTS (SELECT 1 FROM sync.quarantine)")
.fetch_one(&mut *conn)
.await?,
written: 0,
upserts: 0,
deletes: 0,
};
let mut blocked = BTreeMap::new();
self.engine.open_sent(conn).await?;
loop {
let page = self.engine.next(conn, "sync_sent", PAGE).await?;
if page.is_empty() {
break;
}
let mut rest = page.as_slice();
while let Some((i, _, _)) = rest.first() {
let chunk_of = |limit: usize| {
let mut keys: Vec<String> = Vec::new();
let end = rest
.iter()
.position(|(t, key, _)| {
if t != i {
return true;
}
if keys.last() != Some(key) {
if keys.len() == limit {
return true;
}
keys.push(key.clone());
}
false
})
.unwrap_or(rest.len());
(keys, end)
};
let (mut keys, mut end) = chunk_of(RECORD_CHUNK);
let fit = self.engine.fitting(conn, *i, &keys, self.buffered).await?;
if fit < keys.len() {
(keys, end) = chunk_of(fit);
}
let (chunk, after) = rest.split_at(end);
rest = after;
let encoded = self
.engine
.encode(conn, self.codecs, *i, &keys, self.limits.largest_row())
.await?;
for (i, key, replica) in chunk {
if let Some(why) = encoded.failed.get(key) {
blocked.insert((*replica, *i, key.clone()), why.clone());
} else if blocked.is_empty() {
let table = self.engine.table(*i);
writer.push(*replica, table, Some(encoded.rows[key].clone()), key);
}
}
if blocked.is_empty() {
writer.flush(conn, false).await?;
}
}
}
if !blocked.is_empty() {
return Ok(Err(blocked));
}
self.engine.open_dropped(conn).await?;
loop {
let page = self.engine.next(conn, "sync_dropped", PAGE).await?;
if page.is_empty() {
break;
}
for (i, key, replica) in &page {
writer.push(*replica, self.engine.table(*i), None, key);
}
writer.flush(conn, false).await?;
}
for codec in self
.codecs
.iter()
.filter(|c| c.table.mirror == crate::row::Mirror::Setting)
{
for replica in &snapshots {
for line in codec.versions(conn, Some(*replica)).await? {
let key = line.split_once('\t').map_or(line.as_str(), |(key, _)| key);
writer.record(*replica, codec.table.name, key);
}
}
}
writer.end(&self.engine.snapshot_tables());
writer.flush(conn, true).await?;
let replicas: Vec<Uuid> = writer.last.keys().copied().collect();
let (lasts, nonces): (Vec<i64>, Vec<i64>) = writer.last.values().copied().unzip();
let of = |map: &BTreeMap<Uuid, i64>| -> Vec<Option<i64>> {
replicas.iter().map(|r| map.get(r).copied()).collect()
};
sqlx::query(
"INSERT INTO sync.chains AS c \
(replica_id, last_sequence, last_nonce, snapshot_sequence, snapshot_end) \
SELECT * FROM unnest($1::uuid[], $2::bigint[], $3::bigint[], $4::bigint[], $5::bigint[]) \
ON CONFLICT (replica_id) DO UPDATE SET last_sequence = EXCLUDED.last_sequence, \
last_nonce = EXCLUDED.last_nonce, \
snapshot_sequence = coalesce(EXCLUDED.snapshot_sequence, c.snapshot_sequence), \
snapshot_end = coalesce(EXCLUDED.snapshot_end, c.snapshot_end)",
)
.bind(&replicas)
.bind(&lasts)
.bind(&nonces)
.bind(of(&writer.began))
.bind(of(&writer.ended))
.execute(&mut *conn)
.await?;
Ok(Ok((writer.written, writer.upserts, writer.deletes)))
}
async fn bound(&self, conn: &mut Conn, replicas: &[Uuid]) -> Result<Vec<Uuid>, Error> {
let over: Vec<Uuid> = sqlx::query_scalar(
"SELECT o.replica_id FROM sync.outbox o LEFT JOIN sync.bodies b ON b.id = o.body \
WHERE o.replica_id = ANY($1) AND NOT o.snapshot GROUP BY o.replica_id \
HAVING count(*) > $2 \
OR sum(length(o.head) + coalesce(length(b.bytes), 0) + length(o.tail)) > $3",
)
.bind(replicas)
.bind(self.bound.messages)
.bind(self.bound.bytes)
.fetch_all(&mut *conn)
.await?;
for replica in &over {
collapse(
conn,
*replica,
"its outbox outgrew its bound while it could not be reached: provisioned again",
)
.await?;
}
Ok(over)
}
pub(crate) async fn publish<T: Transport>(
&self,
conn: &mut Conn,
transport: &T,
) -> Result<Published, Error> {
let epoch: Uuid = sqlx::query_scalar("SELECT epoch FROM sync.epoch")
.fetch_one(&mut *conn)
.await?;
let due: Vec<Uuid> = sqlx::query_scalar(
"SELECT DISTINCT o.replica_id FROM sync.outbox o \
LEFT JOIN sync.replicas r USING (replica_id) \
WHERE r.retry_after IS NULL OR r.retry_after <= now() ORDER BY 1",
)
.fetch_all(&mut *conn)
.await?;
let conn = tokio::sync::Mutex::new(conn);
let jobs = due.into_iter().map(|replica| {
let job: Job<'_, _> = Box::pin(self.publish_to(&conn, transport, replica, epoch));
job
});
let mut published = Published::default();
for done in at_most(self.concurrency, jobs).await {
let (sent, failing) = done?;
published.sent += sent;
published.failing.extend(failing);
}
Ok(published)
}
async fn publish_to<T: Transport>(
&self,
conn: &tokio::sync::Mutex<&mut Conn>,
transport: &T,
replica: Uuid,
epoch: Uuid,
) -> Result<(usize, Option<(Uuid, String)>), Error> {
let (sent, failure) = self.send(conn, transport, replica, epoch).await?;
let mut conn = conn.lock().await;
let conn: &mut Conn = &mut conn;
let Some(e) = failure else {
sqlx::query(
"INSERT INTO sync.replicas AS r (replica_id, last_sent_at) VALUES ($1, now()) \
ON CONFLICT (replica_id) DO UPDATE SET last_sent_at = now(), \
failures = 0, retry_after = NULL, last_error = NULL, refusals = 0",
)
.bind(replica)
.execute(&mut *conn)
.await?;
return Ok((sent, None));
};
let permanent = !e.is_transient();
let why = if permanent {
format!("the transport refuses a message for good: {e}")
} else {
e.to_string()
};
sqlx::query(
"INSERT INTO sync.replicas (replica_id) VALUES ($1) \
ON CONFLICT (replica_id) DO NOTHING",
)
.bind(replica)
.execute(&mut *conn)
.await?;
let refusals: i32 = sqlx::query_scalar(
"UPDATE sync.replicas SET \
last_sent_at = CASE WHEN $2 THEN now() ELSE last_sent_at END, \
failures = failures + 1, last_error = $3, \
refusals = CASE WHEN $4 THEN refusals + 1 ELSE refusals END, \
retry_after = now() + make_interval(secs => \
least($5 * power(2, least(failures, 32)), $6)) \
WHERE replica_id = $1 RETURNING refusals",
)
.bind(replica)
.bind(sent > 0)
.bind(stored(&why))
.bind(permanent)
.bind(self.backoff.first.as_secs_f64())
.bind(self.backoff.max.as_secs_f64())
.fetch_one(&mut *conn)
.await?;
if refusals >= REFUSALS {
let mut tx = sqlx::Connection::begin(&mut *conn).await?;
sqlx::query("UPDATE sync.replicas SET refusals = 0 WHERE replica_id = $1")
.bind(replica)
.execute(&mut *tx)
.await?;
collapse(
&mut tx,
replica,
&format!("the transport refused its next message for good {refusals} times: provisioned again — {why}"),
)
.await?;
tx.commit().await?;
}
Ok((sent, Some((replica, why))))
}
async fn send<T: Transport>(
&self,
conn: &tokio::sync::Mutex<&mut Conn>,
transport: &T,
replica: Uuid,
epoch: Uuid,
) -> Result<(usize, Option<Error>), Error> {
let mut sent = 0;
let mut after = 0_i64;
loop {
let page: Vec<(i64, i64, Vec<u8>)> = sqlx::query_as(
"SELECT sequence, nonce, message FROM ( \
SELECT o.sequence, o.nonce, o.head || coalesce(b.bytes, '') || o.tail AS message, \
sum(length(o.head) + coalesce(length(b.bytes), 0) + length(o.tail)) \
OVER (ORDER BY o.sequence) AS through \
FROM sync.outbox o LEFT JOIN sync.bodies b ON b.id = o.body \
WHERE o.replica_id = $1 AND o.sequence > $2 \
ORDER BY o.sequence LIMIT 64) AS page \
WHERE through - length(message) < $3 ORDER BY sequence",
)
.bind(replica)
.bind(after)
.bind(i64::try_from(SENT_AT_ONCE).unwrap_or(i64::MAX))
.fetch_all(&mut **conn.lock().await)
.await?;
if page.is_empty() {
return Ok((sent, None));
}
let mut failure = None;
let start = after;
for (sequence, nonce, message) in page {
let id = format!("{epoch}-{sequence}-{nonce}");
let went =
tokio::time::timeout(self.timeout, transport.send(replica, message, &id))
.await
.unwrap_or_else(|_| {
Err(Error::transport(
true,
format!("no answer within {:?}", self.timeout),
))
});
if let Err(e) = went {
failure = Some(e);
break;
}
after = sequence;
sent += 1;
}
if after > start {
drop_outbox(&mut **conn.lock().await, &[replica], after, true).await?;
}
if failure.is_some() {
return Ok((sent, failure));
}
}
}
}
type Job<'a, T> = std::pin::Pin<Box<dyn Future<Output = T> + Send + 'a>>;
async fn at_most<'a, T>(limit: usize, jobs: impl IntoIterator<Item = Job<'a, T>>) -> Vec<T> {
let mut jobs = jobs.into_iter();
let mut running: Vec<Job<'a, T>> = Vec::new();
let mut out = Vec::new();
loop {
running.extend(jobs.by_ref().take(limit.max(1) - running.len()));
if running.is_empty() {
return out;
}
std::future::poll_fn(|cx| {
let before = out.len();
running.retain_mut(|job| match job.as_mut().poll(cx) {
std::task::Poll::Ready(done) => {
out.push(done);
false
}
std::task::Poll::Pending => true,
});
if out.len() > before {
std::task::Poll::Ready(())
} else {
std::task::Poll::Pending
}
})
.await;
}
}
impl Writer {
fn push(&mut self, replica: Uuid, table: &str, bytes: Option<Vec<u8>>, key: &str) {
let wire = envelope::Key {
parts: key.split('|').map(str::to_owned).collect(),
};
self.buffered += bytes.as_ref().map_or(0, Vec::len);
let packer = self
.packers
.get_mut(&replica)
.expect("a packer per replica");
for done in packer.push(table, bytes, wire) {
self.ready(replica, Body::Down(done));
}
self.stuck = self.stuck.min(self.buffered);
if self.buffered > self.room + self.stuck {
let replicas: Vec<Uuid> = self.packers.keys().copied().collect();
for replica in replicas {
if let Some(done) = self.packers.get_mut(&replica).and_then(Packer::relieve) {
self.ready(replica, Body::Down(done));
}
}
self.stuck = self.buffered;
}
}
fn record(&mut self, replica: Uuid, table: &str, key: &str) {
let wire = envelope::Key {
parts: key.split('|').map(str::to_owned).collect(),
};
let packer = self
.packers
.get_mut(&replica)
.expect("a packer per replica");
for done in packer.record(table, wire) {
self.ready(replica, Body::Down(done));
}
}
fn ready(&mut self, replica: Uuid, body: Body) {
if let Body::Down(down) = &body {
self.rows += down
.tables
.iter()
.map(|t| t.upserts.len() + t.deletes.len())
.sum::<usize>();
let carried: usize = down
.tables
.iter()
.flat_map(|t| &t.upserts)
.map(Vec::len)
.sum();
self.buffered = self.buffered.saturating_sub(carried);
self.ready_bytes += carried;
}
self.ready.push(Ready { replica, body });
}
fn end(&mut self, tables: &[&'static str]) {
let replicas: Vec<Uuid> = self.packers.keys().copied().collect();
for replica in replicas {
let packer = self.packers.get_mut(&replica).expect("its packer");
if let Some(done) = packer.finish() {
self.ready(replica, Body::Down(done));
}
if let Some(snapshot_id) = self.snapshots.get(&replica).cloned() {
self.ready(
replica,
Body::EndSnapshot(EndSnapshot {
snapshot_id,
tables: tables.iter().map(|t| (*t).to_owned()).collect(),
..EndSnapshot::default()
}),
);
}
}
}
async fn flush(&mut self, conn: &mut Conn, all: bool) -> Result<(), Error> {
if self.ready.is_empty()
|| (!all && self.rows < RECORD_CHUNK && self.ready_bytes < self.room)
{
return Ok(());
}
let ready = std::mem::take(&mut self.ready);
(self.rows, self.ready_bytes) = (0, 0);
let count = i64::try_from(ready.len())
.map_err(|_| Error::Misuse("more messages than a sequence holds".into()))?;
let mut drawn: Vec<i64> =
sqlx::query_scalar("SELECT nextval('sync.messages') FROM generate_series(1, $1)")
.bind(count)
.fetch_all(&mut *conn)
.await?;
drawn.sort_unstable();
let mut outbox = Outboxed::default();
let mut held = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
let mut gone = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
for (Ready { replica, mut body }, sequence) in ready.into_iter().zip(drawn) {
let number = u64::try_from(sequence)
.map_err(|_| Error::Misuse(format!("sync.messages drew {sequence}")))?;
let (previous, previous_nonce) = self.last.get(&replica).map_or((0, 0), |(s, n)| {
(
u64::try_from(*s).unwrap_or_default(),
u64::try_from(*n).unwrap_or_default(),
)
});
let nonce = envelope::nonce();
Place {
sequence: number,
epoch: self.epoch.to_string(),
nonce,
previous: (previous, previous_nonce),
mac: envelope::seal(&self.secret, replica, self.epoch, number, nonce),
}
.stamp(&mut body);
match &body {
Body::Down(down) => {
for table in &down.tables {
for key in &table.keys {
held.0.push(replica);
held.1.push(table.table.clone());
held.2.push(key.parts.join("|"));
held.3.push(sequence);
}
for key in &table.deletes {
gone.0.push(replica);
gone.1.push(table.table.clone());
gone.2.push(key.parts.join("|"));
gone.3.push(sequence);
}
}
if !down.snapshot_id.is_empty() {
self.began.entry(replica).or_insert(sequence);
}
}
Body::EndSnapshot(_) => {
self.began.entry(replica).or_insert(sequence);
self.ended.insert(replica, sequence);
}
Body::Filed(_) => {}
}
let nonce = i64::try_from(nonce).expect("a nonce is below 2^63");
self.last.insert(replica, (sequence, nonce));
self.written += 1;
outbox.push(replica, sequence, nonce, body);
}
self.upserts += held.0.len();
self.deletes += gone.0.len();
outbox.write(conn).await?;
sqlx::query(
"WITH u AS (SELECT * FROM unnest($1::uuid[], $2::text[], $3::text[], $4::bigint[]) \
AS u(r, t, k, s)), \
n AS MATERIALIZED (SELECT t, k, (SELECT max(id) FROM sync.notes n \
WHERE n.kind = 'row' AND n.relation = r.t AND n.row_key = r.k) AS noted \
FROM (SELECT DISTINCT t, k FROM u) AS r) \
INSERT INTO sync.held (replica_id, relation, row_key, noted, sequence) \
SELECT u.r, u.t, u.k, coalesce(n.noted, 0), u.s FROM u JOIN n USING (t, k) \
ON CONFLICT (replica_id, relation, row_key) DO UPDATE SET \
sent_at = now(), noted = EXCLUDED.noted, sequence = EXCLUDED.sequence",
)
.bind(&held.0)
.bind(&held.1)
.bind(&held.2)
.bind(&held.3)
.execute(&mut *conn)
.await?;
for carried in [&held, &gone].into_iter().filter(|_| self.quarantined) {
sqlx::query(
"UPDATE sync.quarantine q SET carried = u.s \
FROM unnest($1::uuid[], $2::text[], $3::text[], $4::bigint[]) AS u(r, t, k, s) \
WHERE q.replica_id = u.r AND q.relation = u.t AND q.row_key = u.k",
)
.bind(&carried.0)
.bind(&carried.1)
.bind(&carried.2)
.bind(&carried.3)
.execute(&mut *conn)
.await?;
}
sqlx::query(
"WITH u AS (SELECT * FROM unnest($1::uuid[], $2::text[], $3::text[], $4::bigint[]) \
AS u(r, t, k, s)), \
d AS (DELETE FROM sync.held h USING u \
WHERE h.replica_id = u.r AND h.relation = u.t AND h.row_key = u.k \
RETURNING h.replica_id, h.relation, h.row_key, h.noted) \
INSERT INTO sync.gone AS g (replica_id, relation, row_key, sequence, noted) \
SELECT u.r, u.t, u.k, u.s, coalesce(d.noted, 0) FROM u \
LEFT JOIN d ON d.replica_id = u.r AND d.relation = u.t AND d.row_key = u.k \
ON CONFLICT (replica_id, relation, row_key) DO UPDATE SET \
sequence = EXCLUDED.sequence, \
noted = CASE WHEN EXCLUDED.noted = 0 THEN g.noted ELSE EXCLUDED.noted END",
)
.bind(&gone.0)
.bind(&gone.1)
.bind(&gone.2)
.bind(&gone.3)
.execute(&mut *conn)
.await?;
Ok(())
}
}
#[derive(Default)]
struct Outboxed {
replicas: Vec<Uuid>,
sequences: Vec<i64>,
nonces: Vec<i64>,
heads: Vec<Vec<u8>>,
carries: Vec<Option<usize>>,
tails: Vec<Vec<u8>>,
snapshots: Vec<bool>,
bodies: Vec<Vec<u8>>,
found: std::collections::HashMap<u64, Vec<usize>>,
}
impl Outboxed {
fn push(&mut self, replica: Uuid, sequence: i64, nonce: i64, body: Body) {
use std::hash::{Hash as _, Hasher as _};
self.snapshots.push(match &body {
Body::Down(down) => !down.snapshot_id.is_empty(),
_ => true,
});
let envelope::Parts { head, body, tail } = body.parts();
let carries = (!body.is_empty()).then(|| {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
body.hash(&mut hasher);
let same = self.found.entry(hasher.finish()).or_default();
if let Some(&at) = same.iter().find(|&&at| self.bodies[at] == body) {
return at;
}
same.push(self.bodies.len());
self.bodies.push(body);
self.bodies.len() - 1
});
self.replicas.push(replica);
self.sequences.push(sequence);
self.nonces.push(nonce);
self.heads.push(head);
self.carries.push(carries);
self.tails.push(tail);
}
async fn write(self, conn: &mut Conn) -> Result<(), Error> {
let count = i64::try_from(self.bodies.len())
.map_err(|_| Error::Misuse("more bodies than a sequence holds".into()))?;
let ids: Vec<i64> = sqlx::query_scalar(
"SELECT nextval(pg_get_serial_sequence('sync.bodies', 'id')) FROM generate_series(1, $1)",
)
.bind(count)
.fetch_all(&mut *conn)
.await?;
sqlx::query(
"INSERT INTO sync.bodies (id, bytes) SELECT * FROM unnest($1::int8[], $2::bytea[])",
)
.bind(&ids)
.bind(&self.bodies)
.execute(&mut *conn)
.await?;
let carries: Vec<Option<i64>> =
self.carries.iter().map(|at| at.map(|at| ids[at])).collect();
sqlx::query(
"INSERT INTO sync.outbox (replica_id, sequence, nonce, head, body, tail, snapshot) \
SELECT * FROM unnest($1::uuid[], $2::int8[], $3::int8[], $4::bytea[], $5::int8[], \
$6::bytea[], $7::boolean[])",
)
.bind(&self.replicas)
.bind(&self.sequences)
.bind(&self.nonces)
.bind(&self.heads)
.bind(&carries)
.bind(&self.tails)
.bind(&self.snapshots)
.execute(&mut *conn)
.await?;
Ok(())
}
}
pub(crate) async fn drop_outbox(
conn: &mut Conn,
replicas: &[Uuid],
through: i64,
sent: bool,
) -> Result<(), Error> {
sqlx::query(
"WITH gone AS (DELETE FROM sync.outbox WHERE replica_id = ANY($1) AND sequence <= $2 \
RETURNING body), \
sent AS (INSERT INTO sync.replicas (replica_id, last_sent_at) \
SELECT r, now() FROM unnest($1::uuid[]) AS r WHERE $3 \
ON CONFLICT (replica_id) DO UPDATE SET last_sent_at = now()) \
DELETE FROM sync.bodies b WHERE b.id IN (SELECT body FROM gone) \
AND NOT EXISTS (SELECT 1 FROM sync.outbox o WHERE o.body = b.id \
AND NOT (o.replica_id = ANY($1) AND o.sequence <= $2))",
)
.bind(replicas)
.bind(through)
.bind(sent)
.execute(&mut *conn)
.await?;
Ok(())
}
pub(crate) async fn sweep_bodies(conn: &mut Conn) -> Result<(), Error> {
sqlx::query(
"DELETE FROM sync.bodies b WHERE NOT EXISTS (SELECT 1 FROM sync.outbox o WHERE o.body = b.id)",
)
.execute(&mut *conn)
.await?;
Ok(())
}
struct Place {
sequence: u64,
epoch: String,
nonce: u64,
previous: (u64, u64),
mac: Vec<u8>,
}
impl Place {
fn stamp(self, body: &mut Body) {
let Place {
sequence,
epoch,
nonce,
previous: (previous, previous_nonce),
mac,
} = self;
match body {
Body::Down(m) => {
(m.sequence, m.epoch, m.nonce, m.mac) = (sequence, epoch, nonce, mac);
(m.previous, m.previous_nonce) = (previous, previous_nonce);
}
Body::EndSnapshot(m) => {
(m.sequence, m.epoch, m.nonce, m.mac) = (sequence, epoch, nonce, mac);
(m.previous, m.previous_nonce) = (previous, previous_nonce);
}
Body::Filed(_) => {}
}
}
}
async fn collapse(conn: &mut Conn, replica: Uuid, why: &str) -> Result<(), Error> {
tracing::error!(%replica, why, "sync ALARM: a replica's messages dropped — it is provisioned again");
drop_outbox(conn, &[replica], i64::MAX, false).await?;
for sql in [
"DELETE FROM sync.held WHERE replica_id = $1",
"DELETE FROM sync.gone WHERE replica_id = $1",
] {
sqlx::query(sql).bind(replica).execute(&mut *conn).await?;
}
sqlx::query(
"INSERT INTO sync.chains (replica_id, dropped, dropped_at) VALUES ($1, $2, now()) \
ON CONFLICT (replica_id) DO UPDATE SET last_sequence = NULL, last_nonce = NULL, \
snapshot_sequence = NULL, snapshot_end = NULL, dropped = EXCLUDED.dropped, \
dropped_at = now()",
)
.bind(replica)
.bind(stored(why))
.execute(&mut *conn)
.await?;
resync(conn, replica).await?;
Ok(())
}
pub(crate) fn stored(text: &str) -> String {
const STORED: usize = 4 << 10;
let mut end = text.len().min(STORED);
while !text.is_char_boundary(end) {
end -= 1;
}
text[..end].replace('\0', "\u{FFFD}")
}
pub(crate) async fn resync(conn: &mut Conn, replica: Uuid) -> Result<bool, Error> {
let noted = sqlx::query(&crate::plan::resync_sql(
"SELECT $1::uuid, (SELECT partition_id FROM sync.replicas WHERE replica_id = $1)",
))
.bind(replica)
.execute(&mut *conn)
.await?;
Ok(noted.rows_affected() > 0)
}
async fn quarantine(conn: &mut Conn, unsent: &[Unsent]) -> Result<(), Error> {
let Some(first) = unsent.first() else {
return Ok(());
};
tracing::error!(
rows = unsent.len(),
replica = %first.replica,
table = first.table,
key = first.key,
reason = first.reason,
"sync: rows quarantined for replicas, the first of them named — fix the cause and they go on the next retry"
);
let mut rows = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
for row in unsent {
rows.0.push(row.replica);
rows.1.push(row.table);
rows.2.push(row.key.as_str());
rows.3.push(stored(&row.reason));
}
sqlx::query(
"INSERT INTO sync.quarantine AS q (replica_id, relation, row_key, cause, reason) \
SELECT r, t, k, 'server', w FROM unnest($1::uuid[], $2::text[], $3::text[], $4::text[]) AS u(r, t, k, w) \
ON CONFLICT (replica_id, relation, row_key) DO UPDATE SET cause = EXCLUDED.cause, \
reason = EXCLUDED.reason, last_at = now(), attempts = q.attempts + 1, carried = NULL",
)
.bind(&rows.0)
.bind(&rows.1)
.bind(&rows.2)
.bind(&rows.3)
.execute(&mut *conn)
.await?;
Ok(())
}
pub(crate) async fn unpark(conn: &mut Conn) -> Result<u64, Error> {
let done = sqlx::query("UPDATE sync.notes SET parked_at = NULL WHERE parked_at IS NOT NULL")
.execute(&mut *conn)
.await?;
Ok(done.rows_affected())
}
pub(crate) async fn retry_quarantine(conn: &mut Conn) -> Result<u64, Error> {
let done = sqlx::query(
"INSERT INTO sync.notes (kind, relation, row_key) \
SELECT DISTINCT 'retry', q.relation, q.row_key FROM sync.quarantine q \
WHERE NOT EXISTS (SELECT 1 FROM sync.notes n \
WHERE n.done_at IS NULL AND n.kind = 'retry' \
AND n.relation = q.relation AND n.row_key = q.row_key)",
)
.execute(&mut *conn)
.await?;
Ok(done.rows_affected())
}
pub(crate) const UNCONFIRMED: &str = "(c.last_sequence IS NOT NULL \
AND (r.applied_epoch IS DISTINCT FROM (SELECT epoch FROM sync.epoch) \
OR coalesce(r.applied_sequence, 0) < c.last_sequence))";
pub(crate) async fn heal(conn: &mut Conn, engine: &Engine, after: Duration) -> Result<u64, Error> {
let healed: i64 = sqlx::query_scalar(&format!(
"WITH silent AS ( \
SELECT r.replica_id AS id, s.partition FROM sync.replicas r \
JOIN sync.chains c USING (replica_id) \
JOIN {} ON s.id = r.replica_id \
WHERE {UNCONFIRMED} AND r.failures = 0 \
AND NOT EXISTS (SELECT 1 FROM sync.outbox o WHERE o.replica_id = r.replica_id) \
AND greatest(r.last_sent_at, r.applied_at) \
< now() - make_interval(secs => $1 * power(2, least(r.heals, 5)))), \
noted AS ({} RETURNING replica_id), \
counted AS (UPDATE sync.replicas r SET heals = r.heals + 1 FROM noted n \
WHERE r.replica_id = n.replica_id) \
SELECT count(*) FROM noted",
engine.replicas_sql(),
crate::plan::resync_sql("SELECT id, partition FROM silent")
))
.bind(after.as_secs_f64())
.fetch_one(&mut *conn)
.await?;
Ok(u64::try_from(healed).unwrap_or_default())
}
#[cfg(test)]
#[path = "relay_tests.rs"]
mod tests;