use std::time::Duration;
use reddb_wire::replication::CausalBookmark;
use reddb_wire::topology::{Endpoint, ReplicaInfo};
use crate::topology::ClusterMembership;
pub type BookmarkTarget = CausalBookmark;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct WriteEndpoint {
pub addr: String,
pub region: String,
pub term: u64,
}
impl WriteEndpoint {
pub fn serves_term(&self, term: u64) -> bool {
self.term == term
}
pub fn endpoint(&self) -> Endpoint {
Endpoint {
addr: self.addr.clone(),
region: self.region.clone(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ReadEndpoint {
pub addr: String,
pub region: String,
pub healthy: bool,
pub lag_ms: u32,
pub frontier_lsn: u64,
pub rebootstrapping: bool,
pub bookmark_eligible: bool,
}
impl ReadEndpoint {
fn from_replica(info: &ReplicaInfo, bookmark: BookmarkTarget) -> Self {
let eligible =
info.healthy && !info.rebootstrapping && info.last_applied_lsn >= bookmark.commit_lsn();
Self {
addr: info.addr.clone(),
region: info.region.clone(),
healthy: info.healthy,
lag_ms: info.lag_ms,
frontier_lsn: info.last_applied_lsn,
rebootstrapping: info.rebootstrapping,
bookmark_eligible: eligible,
}
}
fn endpoint(&self) -> Endpoint {
Endpoint {
addr: self.addr.clone(),
region: self.region.clone(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RouteKind {
StrongPrimary,
EligibleTarget,
CaughtUpTarget,
FallbackReplica,
FallbackPrimary,
BoundedStalenessReplica,
LocalReplica,
}
impl RouteKind {
pub fn is_fallback(self) -> bool {
matches!(self, Self::FallbackReplica | Self::FallbackPrimary)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RouteDecision {
pub endpoint: Endpoint,
pub kind: RouteKind,
pub waited: Duration,
}
pub trait BookmarkWaiter {
fn elapsed(&self) -> Duration;
fn poll(&mut self, target_addr: &str) -> u64;
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CausalReadOptions {
pub preferred_region: Option<String>,
pub deadline: Duration,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ReadConsistency {
Strong,
Causal {
bookmark: BookmarkTarget,
opts: CausalReadOptions,
},
BoundedStaleness { max_lag: Duration },
Local,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct QueryOptions {
pub consistency: ReadConsistency,
}
impl QueryOptions {
pub fn strong() -> Self {
Self {
consistency: ReadConsistency::Strong,
}
}
pub fn causal(bookmark: BookmarkTarget, opts: CausalReadOptions) -> Self {
Self {
consistency: ReadConsistency::Causal { bookmark, opts },
}
}
pub fn bounded_staleness(max_lag: Duration) -> Self {
Self {
consistency: ReadConsistency::BoundedStaleness { max_lag },
}
}
pub fn local() -> Self {
Self {
consistency: ReadConsistency::Local,
}
}
}
impl Default for QueryOptions {
fn default() -> Self {
Self::local()
}
}
impl CausalReadOptions {
pub fn with_deadline(deadline: Duration) -> Self {
Self {
preferred_region: None,
deadline,
}
}
pub fn prefer_region(mut self, region: impl Into<String>) -> Self {
self.preferred_region = Some(region.into());
self
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoutingTable {
write: WriteEndpoint,
replicas: Vec<ReplicaInfo>,
epoch: u64,
}
impl RoutingTable {
pub fn from_membership(membership: ClusterMembership, term: u64) -> Self {
let ClusterMembership {
primary,
replicas,
epoch,
} = membership;
Self {
write: WriteEndpoint {
addr: primary.addr,
region: primary.region,
term,
},
replicas,
epoch,
}
}
pub fn epoch(&self) -> u64 {
self.epoch
}
pub fn write_endpoint(&self) -> &WriteEndpoint {
&self.write
}
pub fn read_endpoints(&self, bookmark: BookmarkTarget) -> Vec<ReadEndpoint> {
self.replicas
.iter()
.map(|r| ReadEndpoint::from_replica(r, bookmark))
.collect()
}
fn pick_target_index(&self, preferred_region: Option<&str>) -> Option<usize> {
if let Some(region) = preferred_region {
if let Some(i) = self
.replicas
.iter()
.position(|r| r.healthy && !r.rebootstrapping && r.region == region)
{
return Some(i);
}
}
self.replicas
.iter()
.position(|r| r.healthy && !r.rebootstrapping)
}
pub fn route_causal_read(
&self,
bookmark: BookmarkTarget,
opts: &CausalReadOptions,
waiter: &mut dyn BookmarkWaiter,
) -> RouteDecision {
let target_idx = self.pick_target_index(opts.preferred_region.as_deref());
if let Some(idx) = target_idx {
let target = &self.replicas[idx];
if target.last_applied_lsn >= bookmark.commit_lsn() {
return RouteDecision {
endpoint: Endpoint {
addr: target.addr.clone(),
region: target.region.clone(),
},
kind: RouteKind::EligibleTarget,
waited: Duration::ZERO,
};
}
let addr = target.addr.clone();
let region = target.region.clone();
while waiter.elapsed() < opts.deadline {
let frontier = waiter.poll(&addr);
if frontier >= bookmark.commit_lsn() {
return RouteDecision {
endpoint: Endpoint { addr, region },
kind: RouteKind::CaughtUpTarget,
waited: waiter.elapsed(),
};
}
}
return self.fall_back(bookmark, Some(idx), waiter.elapsed());
}
self.fall_back(bookmark, None, waiter.elapsed())
}
pub fn route_read(
&self,
options: &QueryOptions,
waiter: &mut dyn BookmarkWaiter,
) -> RouteDecision {
match &options.consistency {
ReadConsistency::Strong => RouteDecision {
endpoint: self.write.endpoint(),
kind: RouteKind::StrongPrimary,
waited: Duration::ZERO,
},
ReadConsistency::Causal { bookmark, opts } => {
self.route_causal_read(*bookmark, opts, waiter)
}
ReadConsistency::BoundedStaleness { max_lag } => self.route_bounded_staleness(*max_lag),
ReadConsistency::Local => self.route_local(),
}
}
fn route_bounded_staleness(&self, max_lag: Duration) -> RouteDecision {
let max_lag_ms = max_lag.as_millis().min(u32::MAX as u128) as u32;
match self
.replicas
.iter()
.find(|r| r.healthy && !r.rebootstrapping && r.lag_ms <= max_lag_ms)
{
Some(r) => RouteDecision {
endpoint: Endpoint {
addr: r.addr.clone(),
region: r.region.clone(),
},
kind: RouteKind::BoundedStalenessReplica,
waited: Duration::ZERO,
},
None => RouteDecision {
endpoint: self.write.endpoint(),
kind: RouteKind::FallbackPrimary,
waited: Duration::ZERO,
},
}
}
fn route_local(&self) -> RouteDecision {
match self.replicas.iter().find(|r| r.healthy) {
Some(r) => RouteDecision {
endpoint: Endpoint {
addr: r.addr.clone(),
region: r.region.clone(),
},
kind: RouteKind::LocalReplica,
waited: Duration::ZERO,
},
None => RouteDecision {
endpoint: self.write.endpoint(),
kind: RouteKind::FallbackPrimary,
waited: Duration::ZERO,
},
}
}
fn fall_back(
&self,
bookmark: BookmarkTarget,
exclude: Option<usize>,
waited: Duration,
) -> RouteDecision {
let caught_up = self.replicas.iter().enumerate().find(|(i, r)| {
Some(*i) != exclude
&& r.healthy
&& !r.rebootstrapping
&& r.last_applied_lsn >= bookmark.commit_lsn()
});
match caught_up {
Some((_, r)) => RouteDecision {
endpoint: ReadEndpoint::from_replica(r, bookmark).endpoint(),
kind: RouteKind::FallbackReplica,
waited,
},
None => RouteDecision {
endpoint: self.write.endpoint(),
kind: RouteKind::FallbackPrimary,
waited,
},
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use reddb_wire::topology::Endpoint as WireEndpoint;
fn primary() -> WireEndpoint {
WireEndpoint {
addr: "primary:5050".into(),
region: "us-east-1".into(),
}
}
fn replica(addr: &str, region: &str, healthy: bool, frontier: u64) -> ReplicaInfo {
ReplicaInfo {
addr: addr.into(),
region: region.into(),
healthy,
lag_ms: if healthy { 5 } else { u32::MAX },
last_applied_lsn: frontier,
rebootstrapping: false,
}
}
fn lagging_replica(
addr: &str,
region: &str,
healthy: bool,
frontier: u64,
lag_ms: u32,
) -> ReplicaInfo {
ReplicaInfo {
addr: addr.into(),
region: region.into(),
healthy,
lag_ms,
last_applied_lsn: frontier,
rebootstrapping: false,
}
}
fn rebuilding_replica(addr: &str, region: &str, frontier: u64) -> ReplicaInfo {
ReplicaInfo {
addr: addr.into(),
region: region.into(),
healthy: true,
lag_ms: 5,
last_applied_lsn: frontier,
rebootstrapping: true,
}
}
fn membership(replicas: Vec<ReplicaInfo>) -> ClusterMembership {
ClusterMembership {
primary: primary(),
replicas,
epoch: 3,
}
}
struct ScriptedWaiter {
steps: Vec<u64>,
idx: usize,
tick: Duration,
elapsed: Duration,
polled_addrs: Vec<String>,
}
impl ScriptedWaiter {
fn new(steps: Vec<u64>, tick: Duration) -> Self {
Self {
steps,
idx: 0,
tick,
elapsed: Duration::ZERO,
polled_addrs: Vec::new(),
}
}
}
impl BookmarkWaiter for ScriptedWaiter {
fn elapsed(&self) -> Duration {
self.elapsed
}
fn poll(&mut self, target_addr: &str) -> u64 {
self.polled_addrs.push(target_addr.to_string());
self.elapsed += self.tick;
let v = self
.steps
.get(self.idx)
.copied()
.or_else(|| self.steps.last().copied())
.unwrap_or(0);
self.idx += 1;
v
}
}
#[test]
fn write_endpoint_is_primary_keyed_by_term() {
let table = RoutingTable::from_membership(membership(vec![]), 9);
let w = table.write_endpoint();
assert_eq!(w.addr, "primary:5050");
assert_eq!(w.region, "us-east-1");
assert_eq!(w.term, 9);
assert!(w.serves_term(9));
assert!(!w.serves_term(10));
}
#[test]
fn read_endpoints_carry_frontier_and_eligibility() {
let table = RoutingTable::from_membership(
membership(vec![
replica("r-ahead:5050", "us-east-1", true, 200),
replica("r-behind:5050", "us-east-1", true, 90),
replica("r-down:5050", "us-west-2", false, 500),
]),
1,
);
let bookmark = BookmarkTarget::new(1, 100);
let reads = table.read_endpoints(bookmark);
assert_eq!(reads.len(), 3);
assert_eq!(reads[0].frontier_lsn, 200);
assert!(reads[0].bookmark_eligible);
assert_eq!(reads[1].frontier_lsn, 90);
assert!(!reads[1].bookmark_eligible);
assert!(reads[2].frontier_lsn >= 100);
assert!(!reads[2].healthy);
assert!(!reads[2].bookmark_eligible);
}
#[test]
fn eligibility_boundary_is_inclusive_at_commit_lsn() {
let table =
RoutingTable::from_membership(membership(vec![replica("r:5050", "r1", true, 100)]), 1);
let reads = table.read_endpoints(BookmarkTarget::new(1, 100));
assert!(reads[0].bookmark_eligible);
}
#[test]
fn rebootstrapping_replica_is_never_bookmark_eligible_despite_frontier() {
let table = RoutingTable::from_membership(
membership(vec![rebuilding_replica("rebuild:5050", "us-east-1", 999)]),
1,
);
let reads = table.read_endpoints(BookmarkTarget::new(1, 100));
assert_eq!(reads[0].frontier_lsn, 999);
assert!(reads[0].rebootstrapping);
assert!(
!reads[0].bookmark_eligible,
"a rebuilding node must never be bookmark-eligible"
);
}
#[test]
fn route_skips_rebuilding_node_and_falls_back_to_caught_up_peer() {
let table = RoutingTable::from_membership(
membership(vec![
rebuilding_replica("rebuild:5050", "us-east-1", 999),
replica("caught-up:5050", "us-east-1", true, 300),
]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.endpoint.addr, "caught-up:5050");
assert!(!decision.kind.is_fallback());
assert!(
waiter.polled_addrs.iter().all(|a| a != "rebuild:5050"),
"must never poll a rebuilding node"
);
}
#[test]
fn route_falls_back_to_primary_when_every_replica_is_rebuilding() {
let table = RoutingTable::from_membership(
membership(vec![
rebuilding_replica("rebuild-a:5050", "us-east-1", 999),
rebuilding_replica("rebuild-b:5050", "us-west-2", 999),
]),
4,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(4, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
assert!(
waiter.polled_addrs.is_empty(),
"no rebuilding node should be polled"
);
}
#[test]
fn rebuilding_node_excluded_as_fallback_target() {
let table = RoutingTable::from_membership(
membership(vec![
replica("east-lag:5050", "us-east-1", true, 10),
rebuilding_replica("west-rebuild:5050", "us-west-2", 999),
]),
2,
);
let mut waiter = ScriptedWaiter::new(vec![10, 20, 30], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(2, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(25)).prefer_region("us-east-1"),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
}
#[test]
fn route_picks_eligible_target_without_waiting() {
let table = RoutingTable::from_membership(
membership(vec![replica("r-ok:5050", "us-east-1", true, 150)]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::EligibleTarget);
assert_eq!(decision.endpoint.addr, "r-ok:5050");
assert_eq!(decision.waited, Duration::ZERO);
assert!(waiter.polled_addrs.is_empty(), "must not poll on fast path");
}
#[test]
fn route_prefers_region_for_the_target() {
let table = RoutingTable::from_membership(
membership(vec![
replica("east:5050", "us-east-1", true, 150),
replica("west:5050", "us-west-2", true, 150),
]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500))
.prefer_region("us-west-2"),
&mut waiter,
);
assert_eq!(decision.endpoint.addr, "west:5050");
}
#[test]
fn route_waits_and_routes_to_target_once_it_catches_up() {
let table = RoutingTable::from_membership(
membership(vec![replica("r-lag:5050", "us-east-1", true, 50)]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![60, 80, 100], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::CaughtUpTarget);
assert_eq!(decision.endpoint.addr, "r-lag:5050");
assert_eq!(decision.waited, Duration::from_millis(30));
assert_eq!(waiter.polled_addrs.len(), 3);
}
#[test]
fn route_falls_back_to_caught_up_replica_when_target_stays_behind() {
let table = RoutingTable::from_membership(
membership(vec![
replica("east-lag:5050", "us-east-1", true, 10),
replica("west-ok:5050", "us-west-2", true, 300),
]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![10, 20, 30], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(25)).prefer_region("us-east-1"),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackReplica);
assert_eq!(decision.endpoint.addr, "west-ok:5050");
assert!(decision.waited >= Duration::from_millis(25));
assert!(waiter.polled_addrs.iter().all(|a| a == "east-lag:5050"));
}
#[test]
fn route_falls_back_to_primary_when_no_replica_is_caught_up() {
let table = RoutingTable::from_membership(
membership(vec![
replica("r1:5050", "us-east-1", true, 10),
replica("r2:5050", "us-east-1", true, 20),
]),
7,
);
let mut waiter = ScriptedWaiter::new(vec![10, 20], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(7, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(15)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
assert!(decision.kind.is_fallback());
}
#[test]
fn route_falls_back_to_primary_when_no_replica_is_healthy() {
let table = RoutingTable::from_membership(
membership(vec![replica("r-down:5050", "us-east-1", false, 500)]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
assert!(waiter.polled_addrs.is_empty());
}
#[test]
fn route_falls_back_to_primary_when_no_replicas_advertised() {
let table = RoutingTable::from_membership(membership(vec![]), 1);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(500)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::FallbackPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
}
#[test]
fn lagging_replica_never_errors_always_resolves_an_endpoint() {
let shapes = vec![
vec![],
vec![replica("a:5050", "r1", true, 0)],
vec![replica("a:5050", "r1", false, 0)],
vec![
replica("a:5050", "r1", true, 1),
replica("b:5050", "r2", true, 999),
],
];
for shape in shapes {
let table = RoutingTable::from_membership(membership(shape), 1);
let mut waiter = ScriptedWaiter::new(vec![0], Duration::from_millis(10));
let decision = table.route_causal_read(
BookmarkTarget::new(1, 100),
&CausalReadOptions::with_deadline(Duration::from_millis(20)),
&mut waiter,
);
assert!(!decision.endpoint.addr.is_empty());
}
}
#[test]
fn route_read_strong_uses_primary_without_polling() {
let table = RoutingTable::from_membership(
membership(vec![replica("r-ok:5050", "us-east-1", true, 150)]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![150], Duration::from_millis(10));
let decision = table.route_read(&QueryOptions::strong(), &mut waiter);
assert_eq!(decision.kind, RouteKind::StrongPrimary);
assert_eq!(decision.endpoint.addr, "primary:5050");
assert!(waiter.polled_addrs.is_empty());
}
#[test]
fn route_read_bounded_staleness_respects_declared_lag_bound() {
let table = RoutingTable::from_membership(
membership(vec![
lagging_replica("too-stale:5050", "us-east-1", true, 500, 250),
lagging_replica("fresh-enough:5050", "us-east-1", true, 90, 40),
]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![], Duration::from_millis(10));
let decision = table.route_read(
&QueryOptions::bounded_staleness(Duration::from_millis(50)),
&mut waiter,
);
assert_eq!(decision.kind, RouteKind::BoundedStalenessReplica);
assert_eq!(decision.endpoint.addr, "fresh-enough:5050");
assert!(waiter.polled_addrs.is_empty());
let fallback = table.route_read(
&QueryOptions::bounded_staleness(Duration::from_millis(10)),
&mut waiter,
);
assert_eq!(fallback.kind, RouteKind::FallbackPrimary);
assert_eq!(fallback.endpoint.addr, "primary:5050");
}
#[test]
fn route_read_local_never_blocks_on_freshness() {
let table = RoutingTable::from_membership(
membership(vec![lagging_replica(
"very-stale:5050",
"us-east-1",
true,
0,
u32::MAX,
)]),
1,
);
let mut waiter = ScriptedWaiter::new(vec![0, 0, 0], Duration::from_millis(10));
let decision = table.route_read(&QueryOptions::local(), &mut waiter);
assert_eq!(decision.kind, RouteKind::LocalReplica);
assert_eq!(decision.endpoint.addr, "very-stale:5050");
assert_eq!(decision.waited, Duration::ZERO);
assert!(waiter.polled_addrs.is_empty());
}
}