use std::sync::Arc;
use std::sync::atomic::Ordering::{Acquire, Relaxed, Release};
use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize};
use core::cell::{Cell, RefCell};
use yo_common::lock::Lock;
use yo_common::{Code, Error, Result};
use crate::reply::Out;
use super::args::{self, Args};
use super::clients::{self, Client};
use super::pubsub::Envelope;
use super::{Server, Session};
pub(super) const ID_LEN: usize = 40;
const NO_ID: &[u8; ID_LEN] = b"0000000000000000000000000000000000000000";
pub(super) const BACKLOG_BYTES: usize = 1024 * 1024;
pub(super) struct Replica {
row: Arc<Client>,
port: AtomicU64,
ack: AtomicU64,
ack_ms: AtomicU64,
online: AtomicBool,
}
impl Replica {
fn address(&self) -> (String, u64) {
let text = self.row.text.lock();
let peer = text.peer.clone();
drop(text);
let host = match peer.iter().rposition(|&b| b == b':') {
Some(at) => peer[..at].to_vec(),
None => peer.clone(),
};
(
String::from_utf8_lossy(&host).into_owned(),
self.port.load(Relaxed),
)
}
}
pub(crate) struct Replication {
id: Lock<[u8; ID_LEN]>,
id2: Lock<[u8; ID_LEN]>,
offset: AtomicU64,
second: AtomicI64,
backlog: Lock<Backlog>,
rows: Lock<Vec<Arc<Replica>>>,
live: AtomicUsize,
ports: Lock<Vec<(u64, u64)>>,
frozen: AtomicBool,
building: Lock<()>,
on_db: AtomicI64,
}
impl Default for Replication {
fn default() -> Replication {
Replication {
id: Lock::new(make_id()),
id2: Lock::new(*NO_ID),
offset: AtomicU64::new(0),
second: AtomicI64::new(-1),
backlog: Lock::new(Backlog::default()),
rows: Lock::new(Vec::new()),
live: AtomicUsize::new(0),
ports: Lock::new(Vec::new()),
frozen: AtomicBool::new(false),
building: Lock::new(()),
on_db: AtomicI64::new(-1),
}
}
}
#[derive(Default)]
struct Backlog {
ring: Vec<u8>,
at: usize,
filled: usize,
first: u64,
}
impl Backlog {
fn push(&mut self, bytes: &[u8], upto: u64) {
if self.ring.is_empty() {
self.ring = vec![0; BACKLOG_BYTES];
self.first = upto - bytes.len() as u64;
}
for &b in bytes {
self.ring[self.at] = b;
self.at = (self.at + 1) % BACKLOG_BYTES;
if self.filled < BACKLOG_BYTES {
self.filled += 1;
}
}
self.first = upto - self.filled as u64;
}
fn since(&self, from: u64, upto: u64) -> Option<Vec<u8>> {
if self.filled == 0 || from < self.first || from > upto {
return None;
}
let skip = (from - self.first) as usize;
let want = self.filled - skip;
let start = (self.at + BACKLOG_BYTES - self.filled + skip) % BACKLOG_BYTES;
let mut out = Vec::with_capacity(want);
for i in 0..want {
out.push(self.ring[(start + i) % BACKLOG_BYTES]);
}
Some(out)
}
}
#[cfg(test)]
impl Server {
pub(super) fn pretend_replica(&self) {
self.repl.live.store(1, Relaxed);
}
pub(super) fn stream_since(&self, from: u64) -> (String, u64) {
let upto = self.repl.offset.load(Acquire);
let backlog = self.repl.backlog.lock();
let bytes = backlog.since(from, upto).unwrap_or_default();
(String::from_utf8_lossy(&bytes).into_owned(), upto)
}
}
fn make_id() -> [u8; ID_LEN] {
const HEX: &[u8; 16] = b"0123456789abcdef";
let mut raw = [0u8; ID_LEN / 2];
yo_common::entropy::fill(&mut raw);
let mut id = [0u8; ID_LEN];
for (i, byte) in raw.iter().enumerate() {
id[i * 2] = HEX[usize::from(byte >> 4)];
id[i * 2 + 1] = HEX[usize::from(byte & 15)];
}
id
}
thread_local! {
static INSTEAD: RefCell<Option<Vec<Vec<Vec<u8>>>>> = const { RefCell::new(None) };
}
thread_local! {
static ARMED: Cell<bool> = const { Cell::new(false) };
static REAPED: RefCell<Vec<Vec<Vec<u8>>>> = const { RefCell::new(Vec::new()) };
}
pub(super) fn arm(on: bool) -> bool {
ARMED.replace(on)
}
#[must_use]
pub(crate) fn armed() -> bool {
ARMED.get()
}
pub(crate) fn instead(parts: Vec<Vec<u8>>) {
INSTEAD.with(|cell| {
let mut held = cell.borrow_mut();
held.get_or_insert_with(Vec::new).push(parts);
});
}
pub(crate) fn rewrite(parts: &[&[u8]]) {
instead(parts.iter().map(|part| part.to_vec()).collect());
}
pub(crate) fn nothing() {
INSTEAD.with(|cell| {
let mut held = cell.borrow_mut();
held.get_or_insert_with(Vec::new);
});
}
fn taken() -> Option<Vec<Vec<Vec<u8>>>> {
INSTEAD.with(|cell| cell.borrow_mut().take())
}
impl Server {
#[must_use]
pub(crate) fn replicated(&self) -> bool {
self.repl.live.load(Relaxed) != 0
}
#[must_use]
pub(crate) fn replica_count(&self) -> u64 {
self.repl.live.load(Relaxed) as u64
}
pub(crate) fn repl_id(&self) -> [u8; ID_LEN] {
*self.repl.id.lock()
}
#[must_use]
pub(crate) fn repl_offset(&self) -> u64 {
self.repl.offset.load(Acquire)
}
fn replica_rows(&self) -> Vec<Arc<Replica>> {
let rows = self.repl.rows.lock();
yo_alloc::allow(|| rows.clone())
}
fn note_replica_port(&self, id: u64, port: u64) {
const KEEP: usize = 64;
let mut ports = self.repl.ports.lock();
yo_alloc::allow(|| {
if let Some(row) = ports.iter_mut().find(|(who, _)| *who == id) {
row.1 = port;
return;
}
if ports.len() >= KEEP {
ports.remove(0);
}
ports.push((id, port));
});
}
fn take_replica_port(&self, id: u64) -> u64 {
let mut ports = self.repl.ports.lock();
match ports.iter().position(|(who, _)| *who == id) {
Some(at) => ports.remove(at).1,
None => 0,
}
}
fn take_replica(&self, row: &Arc<Client>) -> Arc<Replica> {
let mut rows = self.repl.rows.lock();
let held = yo_alloc::allow(|| {
let held = Arc::new(Replica {
row: Arc::clone(row),
port: AtomicU64::new(self.take_replica_port(row.id)),
ack: AtomicU64::new(0),
ack_ms: AtomicU64::new(self.clock.now_ms()),
online: AtomicBool::new(true),
});
rows.push(Arc::clone(&held));
held
});
self.repl.live.store(rows.len(), Release);
row.set_flag(clients::REPLICA, true);
self.note_here(row.thread.load(Relaxed), 1);
held
}
pub(crate) fn drop_replica(&self, id: u64) {
let mut rows = self.repl.rows.lock();
let Some(at) = rows.iter().position(|r| r.row.id == id) else {
return;
};
let gone = rows.remove(at);
self.repl.live.store(rows.len(), Release);
drop(rows);
gone.row.set_flag(clients::REPLICA, false);
self.note_here(gone.row.thread.load(Relaxed), -1);
}
#[must_use]
pub(crate) fn frozen(&self) -> bool {
self.repl.frozen.load(Relaxed)
}
pub(crate) fn snapshot_at_an_instant(&self) -> (Vec<u8>, u64) {
let building = self.repl.building.lock();
self.repl.frozen.store(true, Release);
for db in &self.dbs {
for stripe in 0..db.width() {
drop(db.hold_stripe(stripe));
}
}
let offset = self.repl.offset.load(Acquire);
let (image, _skipped) = super::persist::build(self);
self.repl.frozen.store(false, Release);
drop(building);
(image, offset)
}
fn replica_of(&self, id: u64) -> Option<Arc<Replica>> {
let rows = self.repl.rows.lock();
rows.iter().find(|r| r.row.id == id).map(Arc::clone)
}
}
fn render(parts: &[&[u8]]) -> Vec<u8> {
let mut out = Vec::with_capacity(16 + parts.iter().map(|p| p.len() + 16).sum::<usize>());
out.extend_from_slice(b"*");
out.extend_from_slice(parts.len().to_string().as_bytes());
out.extend_from_slice(b"\r\n");
for part in parts {
out.extend_from_slice(b"$");
out.extend_from_slice(part.len().to_string().as_bytes());
out.extend_from_slice(b"\r\n");
out.extend_from_slice(part);
out.extend_from_slice(b"\r\n");
}
out
}
fn emit(server: &Server, bytes: Vec<u8>) {
let shared = {
let mut backlog = server.repl.backlog.lock();
let upto = server.repl.offset.load(Relaxed) + bytes.len() as u64;
yo_alloc::allow(|| backlog.push(&bytes, upto));
server.repl.offset.store(upto, Release);
yo_alloc::allow(|| Arc::new(bytes))
};
for held in server.replica_rows() {
if !held.online.load(Relaxed) {
continue;
}
server.post(
held.row.thread.load(Relaxed),
Envelope::raw(
held.row.conn.load(Relaxed),
held.row.id,
Arc::clone(&shared),
),
);
}
}
fn send(server: &Server, db: usize, parts: &[&[u8]]) {
let mut bytes = Vec::new();
if server.repl.on_db.load(Relaxed) != db as i64 {
let n = db.to_string();
bytes = render(&[b"SELECT", n.as_bytes()]);
server.repl.on_db.store(db as i64, Relaxed);
}
bytes.extend_from_slice(&render(parts));
emit(server, bytes);
}
pub(super) fn feed(server: &Server, db: usize, args: Args<'_>, verbatim: bool) {
let instead = taken();
yo_alloc::allow(|| match instead {
None if !verbatim => {}
None => {
let parts: Vec<&[u8]> = (0..args.len()).map(|i| args.get(i)).collect();
send(server, db, &parts);
}
Some(each) => {
for one in &each {
let parts: Vec<&[u8]> = one.iter().map(Vec::as_slice).collect();
send(server, db, &parts);
}
}
});
}
pub(super) fn served(server: &Server, db: usize) {
let Some(each) = taken() else {
return;
};
yo_alloc::allow(|| {
for one in &each {
let parts: Vec<&[u8]> = one.iter().map(Vec::as_slice).collect();
send(server, db, &parts);
}
});
}
pub(super) fn forget() {
let _ = taken();
}
pub(crate) fn reaped(parts: &[&[u8]]) {
if !ARMED.get() {
return;
}
REAPED.with_borrow_mut(|list| {
yo_alloc::allow(|| list.push(parts.iter().map(|part| part.to_vec()).collect()));
});
}
pub(super) fn swept(server: &Server, db: usize) {
if REAPED.with_borrow(Vec::is_empty) {
return;
}
let each = REAPED.with_borrow_mut(core::mem::take);
yo_alloc::allow(|| {
for one in &each {
let parts: Vec<&[u8]> = one.iter().map(Vec::as_slice).collect();
send(server, db, &parts);
}
});
}
pub(super) fn replconf(
server: &Server,
session: &Session,
args: Args<'_>,
out: &mut Out,
) -> Result<()> {
if args.len() < 3 || args.len().is_multiple_of(2) {
return Err(args::wrong_arity("replconf"));
}
let mut at = 1;
while at < args.len() {
let name = args.get(at).to_ascii_lowercase();
match name.as_slice() {
b"listening-port" => {
let port = args.int(at + 1)?;
server.note_replica_port(session.row().id, port.max(0) as u64);
}
b"ack" => {
let ack = args.int(at + 1).unwrap_or(0);
if let Some(held) = server.replica_of(session.row().id) {
held.ack.store(ack.max(0) as u64, Relaxed);
held.ack_ms.store(server.clock.now_ms(), Relaxed);
}
return Ok(());
}
b"getack" => {
return Err(Error::new(
Code::Invalid,
"REPLCONF GETACK is only supported by a replica",
));
}
_ => {}
}
at += 2;
}
out.ok();
Ok(())
}
pub(super) fn psync(
server: &Server,
session: &mut Session,
args: Args<'_>,
out: &mut Out,
) -> Result<()> {
let partial = if args.name().eq_ignore_ascii_case(b"sync") {
if args.len() != 1 {
return Err(args::wrong_arity("sync"));
}
None
} else {
if args.len() != 3 {
return Err(args::wrong_arity("psync"));
}
let asked = args.get(1);
let from = args.int(2)?;
(asked != b"?" && from >= 1).then(|| (asked.to_vec(), from as u64 - 1))
};
if session.running() {
return Err(Error::new(
Code::Invalid,
"PSYNC isn't allowed for DENY BLOCKING client",
));
}
if let Some((asked, from)) = partial
&& let Some(bytes) = catch_up(server, &asked, from)
{
{
let id = server.repl_id();
out.raw(b"+CONTINUE ");
out.raw(&id);
out.raw(b"\r\n");
out.raw(&bytes);
attach(server, session);
return Ok(());
}
}
full(server, session, out)
}
fn catch_up(server: &Server, asked: &[u8], from: u64) -> Option<Vec<u8>> {
let ours = server.repl_id();
let theirs = *server.repl.id2.lock();
let second = server.repl.second.load(Relaxed);
let matches = asked == ours
|| (asked == theirs && second >= 0 && from <= u64::try_from(second).unwrap_or(0));
if !matches {
return None;
}
let upto = server.repl.offset.load(Acquire);
let backlog = server.repl.backlog.lock();
yo_alloc::allow(|| backlog.since(from, upto))
}
fn full(server: &Server, session: &mut Session, out: &mut Out) -> Result<()> {
let (image, offset) = server.snapshot_at_an_instant();
let id = server.repl_id();
out.raw(b"+FULLRESYNC ");
out.raw(&id);
out.raw(b" ");
out.raw(offset.to_string().as_bytes());
out.raw(b"\r\n");
out.raw(b"$");
out.raw(image.len().to_string().as_bytes());
out.raw(b"\r\n");
out.raw(&image);
attach(server, session);
Ok(())
}
fn attach(server: &Server, session: &mut Session) {
super::forget_session(server, session);
server.take_replica(session.row());
}
pub(super) fn caught_up(server: &Server) -> usize {
let upto = server.repl_offset();
server
.replica_rows()
.iter()
.filter(|r| r.online.load(Relaxed) && r.ack.load(Relaxed) >= upto)
.count()
}
pub(super) fn role(server: &Server, out: &mut Out) {
let rows = server.replica_rows();
out.array(3);
out.bulk(b"master");
out.int(server.repl_offset() as i64);
let at = out.len();
let mut n = 0;
for held in &rows {
if !held.online.load(Relaxed) {
continue;
}
let (host, port) = held.address();
out.array(3);
out.bulk(host.as_bytes());
out.bulk(port.to_string().as_bytes());
out.bulk(held.ack.load(Relaxed).to_string().as_bytes());
n += 1;
}
out.close_array(at, n);
}
pub(super) fn info(server: &Server, s: &mut String) {
use core::fmt::Write as _;
let rows = server.replica_rows();
let id = server.repl_id();
let id2 = *server.repl.id2.lock();
let offset = server.repl_offset();
let now = server.clock.now_ms();
let _ = write!(
s,
"# Replication\r\nrole:master\r\nconnected_slaves:{}\r\n",
rows.len()
);
for (i, held) in rows.iter().enumerate() {
let (host, port) = held.address();
let state = if held.online.load(Relaxed) {
"online"
} else {
"wait_bgsave"
};
let lag = now.saturating_sub(held.ack_ms.load(Relaxed)) / 1000;
let _ = write!(
s,
"slave{i}:ip={host},port={port},state={state},offset={},lag={lag}\r\n",
held.ack.load(Relaxed),
);
}
let _ = write!(
s,
"master_failover_state:no-failover\r\n\
master_replid:{}\r\nmaster_replid2:{}\r\n\
master_repl_offset:{offset}\r\nsecond_repl_offset:{}\r\n",
String::from_utf8_lossy(&id),
String::from_utf8_lossy(&id2),
server.repl.second.load(Relaxed),
);
let (active, first, histlen) = {
let backlog = server.repl.backlog.lock();
(
usize::from(!backlog.ring.is_empty()),
if backlog.ring.is_empty() {
0
} else {
backlog.first + 1
},
backlog.filled,
)
};
let _ = write!(
s,
"repl_backlog_active:{active}\r\nrepl_backlog_size:{BACKLOG_BYTES}\r\n\
repl_backlog_first_byte_offset:{first}\r\nrepl_backlog_histlen:{histlen}\r\n\r\n",
);
}