use std::io;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc::{Receiver, Sender, SyncSender, TryRecvError, channel, sync_channel};
use std::thread::JoinHandle;
use melin_transport_core::replication::catchup::CatchUpPublisher;
#[cfg(any(feature = "dpdk", not(feature = "no-persist")))]
use {
melin_app::AppEvent,
melin_transport_core::fence::FenceState,
melin_transport_core::replication::catchup::{
CatchUpResult, can_catch_up_from_journal, catch_up_from_journal_with,
handoff_pass_from_journal, lineage_origin, preflight_snapshot_transfer,
snapshot_transfer_with,
},
melin_transport_core::replication::protocol::{
encode_hash_mismatch, encode_need_snapshot, encode_stream_start,
},
std::path::PathBuf,
};
use super::validation_worker::{WorkerGone, spawn_unpinned};
const JOIN_FRAMES_IN_FLIGHT: usize = 4;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum JoinRequest {
Replica {
last_sequence: u64,
divergent: bool,
},
HandoffPass { from: u64 },
}
enum JoinMessage {
Frame(Vec<u8>),
Done(io::Result<u64>),
}
struct JoinJob {
request: JoinRequest,
frames: SyncSender<JoinMessage>,
spares: Receiver<Vec<u8>>,
cancel: Arc<AtomicBool>,
}
type JoinFn =
Box<dyn FnMut(&JoinRequest, CatchUpPublisher<'_>, &AtomicBool) -> io::Result<u64> + Send>;
pub(crate) struct JoinWorker {
jobs: Option<Sender<JoinJob>>,
handle: Option<JoinHandle<()>>,
}
impl JoinWorker {
pub(crate) fn spawn<F>(name: String, join: F) -> io::Result<Self>
where
F: FnMut(&JoinRequest, CatchUpPublisher<'_>, &AtomicBool) -> io::Result<u64>
+ Send
+ 'static,
{
let (jobs_tx, jobs_rx) = channel::<JoinJob>();
let join: JoinFn = Box::new(join);
let handle = spawn_unpinned(name, move || worker_loop(jobs_rx, join))?;
Ok(JoinWorker {
jobs: Some(jobs_tx),
handle: Some(handle),
})
}
pub(crate) fn submit(&self, request: JoinRequest) -> Result<JoinStream, WorkerGone> {
let (frames_tx, frames_rx) = sync_channel(JOIN_FRAMES_IN_FLIGHT);
let (spares_tx, spares_rx) = channel();
let cancel = Arc::new(AtomicBool::new(false));
let jobs = self.jobs.as_ref().ok_or(WorkerGone)?;
jobs.send(JoinJob {
request,
frames: frames_tx,
spares: spares_rx,
cancel: Arc::clone(&cancel),
})
.map_err(|_| WorkerGone)?;
Ok(JoinStream {
frames: frames_rx,
spares: spares_tx,
held: None,
cancel,
})
}
}
impl Drop for JoinWorker {
fn drop(&mut self) {
self.jobs = None;
if let Some(handle) = self.handle.take()
&& handle.join().is_err()
{
tracing::error!("replica join worker panicked");
}
}
}
fn worker_loop(jobs: Receiver<JoinJob>, mut join: JoinFn) {
for job in jobs {
let JoinJob {
request,
frames,
spares,
cancel,
} = job;
let mut publish = |bytes: &[u8]| -> io::Result<()> {
if cancel.load(Ordering::Acquire) {
return Err(io::Error::other("join abandoned"));
}
let mut buf = spares
.try_recv()
.unwrap_or_else(|_| Vec::with_capacity(bytes.len()));
buf.clear();
buf.extend_from_slice(bytes);
frames
.send(JoinMessage::Frame(buf))
.map_err(|_| io::Error::other("join abandoned"))
};
let result = join(&request, &mut publish, &cancel);
let result = if cancel.load(Ordering::Acquire) {
Err(io::Error::other("join abandoned"))
} else {
result
};
let _ = frames.send(JoinMessage::Done(result));
}
}
pub(crate) struct JoinStream {
frames: Receiver<JoinMessage>,
spares: Sender<Vec<u8>>,
held: Option<Vec<u8>>,
cancel: Arc<AtomicBool>,
}
#[derive(Debug)]
pub(crate) enum Pump {
Pending { queued: usize },
Done(io::Result<u64>),
WorkerGone,
}
impl JoinStream {
pub(crate) fn pump(
&mut self,
budget: usize,
capacity: usize,
mut queue: impl FnMut(&[u8]) -> bool,
) -> Pump {
let mut queued = 0usize;
loop {
if queued > 0 && queued >= budget {
return Pump::Pending { queued };
}
let frame = match self.held.take() {
Some(frame) => frame,
None => match self.frames.try_recv() {
Ok(JoinMessage::Frame(frame)) => frame,
Ok(JoinMessage::Done(result)) => return Pump::Done(result),
Err(TryRecvError::Empty) => return Pump::Pending { queued },
Err(TryRecvError::Disconnected) => return Pump::WorkerGone,
},
};
if frame.len() > capacity {
return Pump::Done(Err(io::Error::other(format!(
"a {}-byte join frame exceeds the replica socket's {capacity}-byte queue",
frame.len()
))));
}
if !queue(&frame) {
self.held = Some(frame);
return Pump::Pending { queued };
}
queued += frame.len();
let _ = self.spares.send(frame);
}
}
}
impl JoinStream {
pub(crate) fn is_blocked(&self) -> bool {
self.held.is_some()
}
}
impl Drop for JoinStream {
fn drop(&mut self) {
self.cancel.store(true, Ordering::Release);
}
}
#[derive(Debug, Default)]
pub(crate) struct JoinStallWatch {
armed: bool,
since: Option<std::time::Instant>,
}
impl JoinStallWatch {
pub(crate) fn arm(&mut self) {
self.armed = true;
}
pub(crate) fn observe(
&mut self,
moved: bool,
blocked: bool,
now: std::time::Instant,
) -> std::time::Duration {
self.since = match (self.armed, blocked, moved) {
(false, _, _) | (true, false, _) => None,
(true, true, true) => Some(now),
(true, true, false) => Some(self.since.unwrap_or(now)),
};
self.since.map_or(std::time::Duration::ZERO, |since| {
now.saturating_duration_since(since)
})
}
pub(crate) fn restart(&mut self) {
self.since = None;
}
}
#[cfg(any(feature = "dpdk", not(feature = "no-persist")))]
pub(crate) struct JoinContext {
pub(crate) journal_path: PathBuf,
pub(crate) fence_state: Arc<FenceState>,
}
#[cfg(any(feature = "dpdk", not(feature = "no-persist")))]
pub(crate) fn stream_join<E: AppEvent>(
ctx: &JoinContext,
request: &JoinRequest,
publish: CatchUpPublisher<'_>,
cancel: &AtomicBool,
) -> io::Result<u64> {
let journal_path = ctx.journal_path.as_path();
let (last_sequence, divergent) = match *request {
JoinRequest::Replica {
last_sequence,
divergent,
} => (last_sequence, divergent),
JoinRequest::HandoffPass { from } => {
return handoff_pass_from_journal::<E>(journal_path, from, publish, cancel);
}
};
let can_catch_up = !divergent
&& can_catch_up_from_journal(journal_path, last_sequence)
.map_err(|e| io::Error::other(format!("catch-up probe: {e}")))?;
let mut frame = Vec::with_capacity(128);
let result = if can_catch_up {
let origin = lineage_origin(journal_path)?;
encode_stream_start(
last_sequence,
origin.starting_sequence,
origin.anchor_hash,
origin.genesis_entries,
ctx.fence_state.epoch(),
&mut frame,
);
publish(&frame)?;
catch_up_from_journal_with::<E>(journal_path, last_sequence, publish, cancel)?
} else {
if divergent {
encode_hash_mismatch(&mut frame);
} else {
encode_need_snapshot(&mut frame);
}
preflight_snapshot_transfer(journal_path)?;
publish(&frame)?;
snapshot_transfer_with::<E>(journal_path, publish, cancel)?
};
match result {
CatchUpResult::Ok(end) => Ok(end),
CatchUpResult::NeedSnapshot if can_catch_up => Err(io::Error::other(
"catch-up found no starting segment after the probe found one",
)),
CatchUpResult::NeedSnapshot => Err(io::Error::other(
"catch-up failed even after snapshot transfer",
)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Mutex;
use std::time::Duration;
fn frames_job(
count: u8,
) -> impl FnMut(&JoinRequest, CatchUpPublisher<'_>, &AtomicBool) -> io::Result<u64> + Send {
move |_req, publish, _cancel| {
for i in 0..count {
publish(&[i])?;
}
Ok(u64::from(count))
}
}
fn request() -> JoinRequest {
JoinRequest::Replica {
last_sequence: 0,
divergent: false,
}
}
fn pump_to_end(stream: &mut JoinStream) -> (Vec<Vec<u8>>, io::Result<u64>) {
let mut sent = Vec::new();
loop {
match stream.pump(usize::MAX, usize::MAX, |f| {
sent.push(f.to_vec());
true
}) {
Pump::Pending { .. } => std::thread::yield_now(),
Pump::Done(result) => return (sent, result),
Pump::WorkerGone => panic!("the worker died"),
}
}
}
#[test]
fn frames_arrive_in_order_then_the_result() {
let worker = JoinWorker::spawn("test-join".into(), frames_job(20)).unwrap();
let mut stream = worker.submit(request()).unwrap();
let (sent, result) = pump_to_end(&mut stream);
assert_eq!(sent, (0..20u8).map(|i| vec![i]).collect::<Vec<_>>());
assert_eq!(result.unwrap(), 20);
}
#[test]
fn a_pump_never_waits_on_a_stalled_job() {
let (release_tx, release_rx) = channel::<()>();
let gate = Mutex::new(release_rx);
let worker = JoinWorker::spawn(
"test-join".into(),
move |_req: &JoinRequest, publish: CatchUpPublisher<'_>, _c: &AtomicBool| {
publish(b"before")?;
gate.lock()
.expect("test mutex poisoned")
.recv()
.expect("release channel closed");
publish(b"after")?;
Ok(7)
},
)
.unwrap();
let mut stream = worker.submit(request()).unwrap();
let mut sent = Vec::new();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
while sent.is_empty() {
assert!(
std::time::Instant::now() < deadline,
"the first frame never came"
);
let _ = stream.pump(usize::MAX, usize::MAX, |f| {
sent.push(f.to_vec());
true
});
}
assert_eq!(sent, [b"before".to_vec()]);
for _ in 0..100 {
let started = std::time::Instant::now();
assert!(matches!(
stream.pump(usize::MAX, usize::MAX, |_| panic!("nothing is ready")),
Pump::Pending { queued: 0 }
));
assert!(
started.elapsed() < Duration::from_secs(1),
"the pump waited"
);
}
release_tx.send(()).unwrap();
let (rest, result) = pump_to_end(&mut stream);
assert_eq!(rest, [b"after".to_vec()]);
assert_eq!(result.unwrap(), 7);
}
#[test]
fn a_refused_frame_is_resent_first() {
let worker = JoinWorker::spawn("test-join".into(), frames_job(6)).unwrap();
let mut stream = worker.submit(request()).unwrap();
let mut sent = Vec::new();
let mut accept = false;
let result = loop {
match stream.pump(usize::MAX, usize::MAX, |f| {
accept = !accept;
if accept {
sent.push(f.to_vec());
}
accept
}) {
Pump::Pending { .. } => std::thread::yield_now(),
Pump::Done(result) => break result,
Pump::WorkerGone => panic!("the worker died"),
}
};
assert_eq!(sent, (0..6u8).map(|i| vec![i]).collect::<Vec<_>>());
assert_eq!(result.unwrap(), 6);
}
#[test]
fn the_budget_bounds_a_pump_but_one_frame_always_moves() {
let worker = JoinWorker::spawn(
"test-join".into(),
|_req: &JoinRequest, publish: CatchUpPublisher<'_>, _c: &AtomicBool| {
for _ in 0..3 {
publish(&[0u8; 100])?;
}
Ok(3)
},
)
.unwrap();
let mut stream = worker.submit(request()).unwrap();
let mut frames_this_pump = Vec::new();
let result = loop {
let mut count = 0;
match stream.pump(10, usize::MAX, |_| {
count += 1;
true
}) {
Pump::Pending { queued } => {
assert!(count <= 1, "a 10-byte budget took {count} 100-byte frames");
assert_eq!(queued, count * 100);
if count > 0 {
frames_this_pump.push(count);
}
std::thread::yield_now();
}
Pump::Done(result) => break result,
Pump::WorkerGone => panic!("the worker died"),
}
};
assert_eq!(frames_this_pump.len(), 3);
assert_eq!(result.unwrap(), 3);
}
#[test]
fn a_frame_larger_than_the_socket_fails_the_join() {
let worker = JoinWorker::spawn(
"test-join".into(),
|_req: &JoinRequest, publish: CatchUpPublisher<'_>, _c: &AtomicBool| {
publish(&[0u8; 100])?;
Ok(1)
},
)
.unwrap();
let mut stream = worker.submit(request()).unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
match stream.pump(usize::MAX, 99, |_| panic!("must not be offered")) {
Pump::Pending { .. } => {
assert!(std::time::Instant::now() < deadline, "the frame never came");
std::thread::yield_now();
}
Pump::Done(result) => {
assert!(result.is_err());
break;
}
Pump::WorkerGone => panic!("the worker died"),
}
}
}
#[test]
fn a_failed_job_reports_its_error() {
let worker = JoinWorker::spawn(
"test-join".into(),
|_req: &JoinRequest, publish: CatchUpPublisher<'_>, _c: &AtomicBool| {
publish(b"verdict")?;
Err(io::Error::new(io::ErrorKind::NotFound, "no snapshot"))
},
)
.unwrap();
let mut stream = worker.submit(request()).unwrap();
let (sent, result) = pump_to_end(&mut stream);
assert_eq!(sent, [b"verdict".to_vec()]);
assert_eq!(result.unwrap_err().kind(), io::ErrorKind::NotFound);
}
#[test]
fn an_abandoned_join_frees_the_worker_for_the_next() {
let worker = JoinWorker::spawn(
"test-join".into(),
|req: &JoinRequest, publish: CatchUpPublisher<'_>, cancel: &AtomicBool| {
let JoinRequest::HandoffPass { from } = *req else {
panic!("this test submits passes only");
};
for _ in 0..1000 {
if cancel.load(Ordering::Acquire) {
break;
}
publish(&from.to_le_bytes())?;
}
Ok(from)
},
)
.unwrap();
let abandoned = worker.submit(JoinRequest::HandoffPass { from: 1 }).unwrap();
std::thread::sleep(Duration::from_millis(50));
drop(abandoned);
let mut stream = worker.submit(JoinRequest::HandoffPass { from: 2 }).unwrap();
let (sent, result) = pump_to_end(&mut stream);
assert_eq!(sent.len(), 1000);
assert!(sent.iter().all(|f| f == &2u64.to_le_bytes()));
assert_eq!(result.unwrap(), 2);
}
#[test]
fn dropping_a_slot_mid_join_does_not_hang() {
let worker = JoinWorker::spawn(
"test-join".into(),
|_req: &JoinRequest, publish: CatchUpPublisher<'_>, _c: &AtomicBool| {
loop {
publish(&[0u8; 8])?;
}
},
)
.unwrap();
let stream = worker.submit(request()).unwrap();
std::thread::sleep(Duration::from_millis(50));
let (done_tx, done_rx) = channel();
std::thread::spawn(move || {
drop(stream);
drop(worker);
done_tx.send(()).unwrap();
});
done_rx
.recv_timeout(Duration::from_secs(10))
.expect("dropping the stream and then the worker must not hang");
}
#[test]
fn a_cancelled_join_is_not_reported_as_done() {
let (seen_tx, seen_rx) = channel::<()>();
let (go_tx, go_rx) = channel::<()>();
let gates = Mutex::new((seen_tx, go_rx));
let worker = JoinWorker::spawn(
"test-join".into(),
move |_req: &JoinRequest, _p: CatchUpPublisher<'_>, _c: &AtomicBool| {
let gates = gates.lock().expect("test mutex poisoned");
gates.0.send(()).expect("test alive");
gates.1.recv().expect("test alive");
Ok(5)
},
)
.unwrap();
let stream = worker.submit(request()).unwrap();
seen_rx.recv().unwrap();
stream.cancel.store(true, Ordering::Release);
go_tx.send(()).unwrap();
match stream.frames.recv().unwrap() {
JoinMessage::Done(result) => assert!(result.is_err()),
JoinMessage::Frame(_) => panic!("no frame was published"),
}
}
#[test]
fn a_panicking_job_is_reported_as_a_dead_worker() {
let worker = JoinWorker::spawn(
"test-join".into(),
|_req: &JoinRequest, _p: CatchUpPublisher<'_>, _c: &AtomicBool| -> io::Result<u64> {
panic!("job bug")
},
)
.unwrap();
let mut stream = worker.submit(request()).unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
match stream.pump(usize::MAX, usize::MAX, |_| true) {
Pump::Pending { .. } => {
assert!(std::time::Instant::now() < deadline, "no outcome");
std::thread::yield_now();
}
Pump::WorkerGone => break,
Pump::Done(r) => panic!("expected a dead worker, got {r:?}"),
}
}
drop(worker);
}
#[cfg(not(feature = "no-persist"))]
mod stream {
use super::*;
use counter_server::CounterEvent;
use melin_journal::{BufferedWriter, JournalEvent, JournalWrite};
use melin_transport_core::replication::protocol::{PrimaryMessage, decode_primary_message};
fn ctx(journal_path: PathBuf) -> JoinContext {
JoinContext {
journal_path,
fence_state: Arc::new(FenceState::new(3)),
}
}
fn journal(dir: &std::path::Path, events: u64) -> PathBuf {
let path = dir.join("primary.journal");
let mut w = BufferedWriter::<CounterEvent>::create(&path).expect("create");
for _ in 0..events {
w.append(&JournalEvent::App(CounterEvent::Increment { amount: 1 }))
.expect("append");
}
drop(w);
path
}
fn run(ctx: &JoinContext, request: JoinRequest) -> (Vec<Vec<u8>>, io::Result<u64>) {
let mut sent = Vec::new();
let never = AtomicBool::new(false);
let result = stream_join::<CounterEvent>(
ctx,
&request,
&mut |frame: &[u8]| {
sent.push(frame.to_vec());
Ok(())
},
&never,
);
(sent, result)
}
fn decode(frame: &[u8]) -> PrimaryMessage {
decode_primary_message(&frame[4..]).expect("a primary message")
}
#[test]
fn a_replica_within_the_journal_gets_stream_start_then_history() {
let dir = tempfile::tempdir().unwrap();
let ctx = ctx(journal(dir.path(), 5));
let (sent, result) = run(
&ctx,
JoinRequest::Replica {
last_sequence: 2,
divergent: false,
},
);
assert_eq!(result.unwrap(), 5);
match decode(&sent[0]) {
PrimaryMessage::StreamStart {
start_sequence,
epoch,
..
} => {
assert_eq!((start_sequence, epoch), (2, 3));
}
other => panic!("expected StreamStart first, got {other:?}"),
}
assert!(sent.len() > 1, "the history follows the StreamStart");
}
#[test]
fn a_divergent_replica_gets_no_verdict_without_a_snapshot() {
let dir = tempfile::tempdir().unwrap();
let ctx = ctx(journal(dir.path(), 3));
let (sent, result) = run(
&ctx,
JoinRequest::Replica {
last_sequence: 2,
divergent: true,
},
);
assert!(result.is_err());
assert!(
sent.is_empty(),
"nothing may be sent: {} frames",
sent.len()
);
}
#[test]
fn a_handoff_pass_streams_only_the_entries_past_its_position() {
let dir = tempfile::tempdir().unwrap();
let ctx = ctx(journal(dir.path(), 5));
let (sent, result) = run(&ctx, JoinRequest::HandoffPass { from: 3 });
assert_eq!(result.unwrap(), 5);
let seqs: Vec<u64> = sent
.iter()
.flat_map(|frame| {
melin_transport_core::replication_wire::try_decode_input_batch::<CounterEvent>(
&frame[4..],
)
.expect("only InputBatch frames")
.into_iter()
.map(|slot| slot.sequence)
})
.collect();
assert_eq!(seqs, [4, 5]);
let (sent, result) = run(&ctx, JoinRequest::HandoffPass { from: 5 });
assert_eq!(result.unwrap(), 5);
assert!(sent.is_empty());
}
}
#[test]
fn the_stall_watch_counts_only_refusals_without_progress() {
let t0 = std::time::Instant::now();
let at = |ms: u64| t0 + Duration::from_millis(ms);
let mut watch = JoinStallWatch::default();
watch.arm();
assert_eq!(watch.observe(true, false, at(0)), Duration::ZERO);
assert_eq!(watch.observe(false, true, at(10)), Duration::ZERO);
assert_eq!(
watch.observe(false, true, at(1_010)),
Duration::from_secs(1)
);
assert_eq!(watch.observe(true, true, at(2_000)), Duration::ZERO);
assert_eq!(
watch.observe(false, true, at(7_000)),
Duration::from_secs(5)
);
assert_eq!(watch.observe(false, false, at(60_000)), Duration::ZERO);
assert_eq!(watch.observe(false, true, at(61_000)), Duration::ZERO);
}
#[test]
fn the_stall_watch_counts_nothing_before_the_replica_acks() {
let t0 = std::time::Instant::now();
let at = |ms: u64| t0 + Duration::from_millis(ms);
let mut watch = JoinStallWatch::default();
assert_eq!(watch.observe(false, true, at(0)), Duration::ZERO);
assert_eq!(watch.observe(false, true, at(600_000)), Duration::ZERO);
watch.arm();
assert_eq!(watch.observe(false, true, at(601_000)), Duration::ZERO);
assert_eq!(
watch.observe(false, true, at(603_000)),
Duration::from_secs(2)
);
watch.arm();
assert_eq!(
watch.observe(false, true, at(604_000)),
Duration::from_secs(3)
);
}
#[test]
fn a_restarted_stall_watch_stays_armed() {
let t0 = std::time::Instant::now();
let at = |ms: u64| t0 + Duration::from_millis(ms);
let mut watch = JoinStallWatch::default();
watch.arm();
assert_eq!(watch.observe(false, true, at(0)), Duration::ZERO);
assert_eq!(
watch.observe(false, true, at(4_000)),
Duration::from_secs(4)
);
watch.restart();
assert_eq!(watch.observe(false, true, at(5_000)), Duration::ZERO);
assert_eq!(
watch.observe(false, true, at(6_000)),
Duration::from_secs(1),
"still armed after the restart"
);
}
#[test]
fn a_held_frame_reads_as_blocked() {
let worker = JoinWorker::spawn("test-join".into(), frames_job(2)).unwrap();
let mut stream = worker.submit(request()).unwrap();
let deadline = std::time::Instant::now() + Duration::from_secs(10);
loop {
assert!(!stream.is_blocked());
match stream.pump(usize::MAX, usize::MAX, |_| false) {
Pump::Pending { queued: 0 } if stream.is_blocked() => break,
Pump::Pending { .. } => {
assert!(std::time::Instant::now() < deadline, "no frame came");
std::thread::yield_now();
}
other => panic!("expected a refused frame, got {other:?}"),
}
}
let (sent, result) = pump_to_end(&mut stream);
assert_eq!(sent.len(), 2);
assert!(!stream.is_blocked());
assert_eq!(result.unwrap(), 2);
}
}