use std::{
collections::{HashMap, hash_map::Entry},
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use serde::Deserialize;
use uuid::Uuid;
use crate::instance::{Instance, State};
static NEXT_IDX: AtomicUsize = AtomicUsize::new(0);
type ConnectionType = String;
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct BucketRangeEntry {
pub tier_name: String,
pub bucket_id_start: u64,
pub bucket_id_end: u64,
pub current_replicaset_uuid: Uuid,
}
pub mod view {
use std::{collections::HashMap, ops::RangeInclusive};
use uuid::Uuid;
use super::ConnectionType;
use crate::error::{Error, IncompleteBucketMap};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Topology {
pub(super) tiers: Vec<Tier>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Tier {
pub(super) name: String,
pub(super) bucket_count: u64,
pub(super) replicasets: Vec<Replicaset>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Replicaset {
pub(super) uuid: Uuid,
pub(super) name: String,
pub(super) master: Option<Instance>,
pub(super) replicas: Vec<Instance>,
pub(super) ranges: Vec<RangeInclusive<u64>>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Instance {
pub(super) name: String,
pub(super) addresses: HashMap<ConnectionType, String>,
pub(super) online: bool,
pub(super) is_master: bool,
pub(super) raft_id: u64,
pub(super) uuid: Uuid,
}
pub type BucketMap<'a> = HashMap<Uuid, &'a [RangeInclusive<u64>]>;
impl Topology {
#[must_use]
pub fn tiers(&self) -> &[Tier] {
&self.tiers
}
}
impl Tier {
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn replicasets(&self) -> &[Replicaset] {
&self.replicasets
}
#[must_use]
pub fn bucket_count(&self) -> u64 {
self.bucket_count
}
pub fn bucket_map(&self) -> Result<BucketMap<'_>, Error> {
let mut ranges: Vec<_> = self.replicasets.iter().flat_map(|rs| &rs.ranges).collect();
ranges.sort_unstable_by_key(|r| *r.start());
let mut covered = 0;
for r in ranges {
if *r.start() != covered + 1 || r.is_empty() {
break;
}
covered = *r.end();
}
if covered != self.bucket_count {
return Err(Error::IncompleteBucketMap(Box::new(IncompleteBucketMap {
tier: self.name.clone(),
covered,
expected: self.bucket_count,
})));
}
Ok(self
.replicasets
.iter()
.map(|rs| (rs.uuid(), rs.ranges()))
.collect())
}
}
impl Replicaset {
#[must_use]
pub fn uuid(&self) -> Uuid {
self.uuid
}
#[must_use]
pub fn master(&self) -> Option<&Instance> {
self.master.as_ref()
}
#[must_use]
pub fn replicas(&self) -> &[Instance] {
&self.replicas
}
pub fn instances(&self) -> impl Iterator<Item = &Instance> {
self.master.iter().chain(&self.replicas)
}
#[must_use]
pub fn ranges(&self) -> &[RangeInclusive<u64>] {
&self.ranges
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
}
impl Instance {
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn address(&self, connection_type: impl AsRef<str>) -> Option<&str> {
self.addresses
.get(connection_type.as_ref())
.map(String::as_str)
}
#[must_use]
pub fn is_online(&self) -> bool {
self.online
}
#[must_use]
pub fn is_master(&self) -> bool {
self.is_master
}
#[must_use]
pub fn raft_id(&self) -> u64 {
self.raft_id
}
#[must_use]
pub fn uuid(&self) -> Uuid {
self.uuid
}
}
}
#[derive(Debug, Default)]
pub(crate) struct Topology {
pub(crate) instances: HashMap<String, Arc<Instance>>,
pgproto_tiers: Vec<Tier>,
other_tiers: HashMap<ConnectionType, Vec<Tier>>,
}
#[derive(Debug, Clone)]
struct Tier {
name: String,
bucket_count: u64,
replicasets: Vec<Replicaset>,
bucket_map: Vec<BucketSlot>,
}
#[derive(Debug, Clone)]
struct BucketSlot {
start: u64,
end: u64,
replicaset_idx: usize,
}
#[derive(Debug, Clone)]
struct Replicaset {
uuid: Uuid,
master: Arc<Instance>,
replicas: Vec<Arc<Instance>>,
}
impl From<(Vec<pg::Row>, Vec<pg::Row>)> for Topology {
fn from((topology_rows, bucket_rows): (Vec<pg::Row>, Vec<pg::Row>)) -> Self {
let mut by_type: HashMap<ConnectionType, Vec<pg::Row>> = HashMap::new();
for row in topology_rows {
let Ok(connection_type) = row.try_get::<_, &str>("connection_type") else {
continue;
};
if let Some(rows) = by_type.get_mut(connection_type) {
rows.push(row);
} else {
by_type.insert(connection_type.to_string(), vec![row]);
}
}
let mut buckets: Vec<BucketRangeEntry> = Vec::new();
for row in bucket_rows {
let entries = if let Ok(json) = row.try_get::<_, serde_json::Value>(0) {
serde_json::from_value::<Vec<BucketRangeEntry>>(json)
} else if let Ok(s) = row.try_get::<_, &str>(0) {
serde_json::from_str(s)
} else {
log::error!("failed to read bucket ranges column");
continue;
};
match entries {
Ok(v) => buckets.extend(v),
Err(e) => log::error!("failed to parse bucket ranges: {e}"),
}
}
let mut other_tiers: HashMap<ConnectionType, Vec<Tier>> = by_type
.into_iter()
.map(|(connection_type, rows)| (connection_type, build_tiers(rows, &buckets)))
.collect();
let pgproto_tiers = other_tiers.remove("pgproto").unwrap_or_default();
let instances = pgproto_tiers
.iter()
.flat_map(|t| &t.replicasets)
.flat_map(|rs| {
rs.master
.is_master
.then_some(&rs.master)
.into_iter()
.chain(&rs.replicas)
})
.map(|inst| (inst.name.clone(), inst.clone()))
.collect();
Topology {
instances,
pgproto_tiers,
other_tiers,
}
}
}
fn build_tiers(topology_rows: Vec<pg::Row>, buckets: &[BucketRangeEntry]) -> Vec<Tier> {
let mut tiers: Vec<Tier> = Vec::new();
for row in topology_rows {
let inst = match Instance::try_from(row) {
Ok(v) => v,
Err(e) => {
log::error!("failed to convert instance: {e}");
continue;
}
};
let inst = Arc::new(inst);
let tier = if let Some(t) = tiers.iter_mut().find(|t| t.name == inst.tier) {
t
} else {
tiers.push(Tier {
name: inst.tier.clone(),
bucket_count: inst.bucket_count,
replicasets: Vec::new(),
bucket_map: Vec::new(),
});
tiers.last_mut().expect("just pushed")
};
let rs = if let Some(r) = tier
.replicasets
.iter_mut()
.find(|r| r.uuid == inst.replicaset_uuid)
{
r
} else {
tier.replicasets.push(Replicaset {
uuid: inst.replicaset_uuid,
master: inst.clone(),
replicas: Vec::new(),
});
tier.replicasets.last_mut().expect("just pushed")
};
if inst.is_master {
rs.master = inst.clone();
} else {
rs.replicas.push(inst.clone());
}
}
for entry in buckets {
let Some(tier) = tiers.iter_mut().find(|t| t.name == entry.tier_name) else {
continue;
};
let Some(rs_idx) = tier
.replicasets
.iter()
.position(|r| r.uuid == entry.current_replicaset_uuid)
else {
continue;
};
tier.bucket_map.push(BucketSlot {
start: entry.bucket_id_start,
end: entry.bucket_id_end,
replicaset_idx: rs_idx,
});
}
for tier in &mut tiers {
tier.bucket_map.sort_unstable_by_key(|s| s.start);
}
tiers
}
impl Topology {
pub(crate) fn new() -> Self {
Self::default()
}
pub(crate) fn merge(&mut self, incoming: Topology) {
for (name, existing) in &self.instances {
if !incoming.instances.contains_key(name) {
existing.set_state(State::Offline);
}
}
for (name, inc_inst) in incoming.instances {
match self.instances.entry(name) {
Entry::Occupied(mut e) => {
inc_inst.set_idx(e.get().idx());
e.insert(inc_inst);
}
Entry::Vacant(e) => {
inc_inst.set_idx(NEXT_IDX.fetch_add(1, Ordering::Relaxed));
e.insert(inc_inst);
}
}
}
self.pgproto_tiers = incoming.pgproto_tiers;
self.other_tiers = incoming.other_tiers;
}
pub(crate) fn resolve(&self, tier: &str, bucket_id: u64, on_master: bool) -> Option<usize> {
let t = self.pgproto_tiers.iter().find(|t| t.name == tier)?;
let rs = t.lookup(bucket_id)?;
if on_master || rs.replicas.is_empty() {
return if rs.master.is_online() {
Some(rs.master.idx())
} else {
None
};
}
let online: Vec<_> = rs.replicas.iter().filter(|r| r.is_online()).collect();
if online.is_empty() {
return if rs.master.is_online() {
Some(rs.master.idx())
} else {
None
};
}
let idx =
usize::try_from(bucket_id % online.len() as u64).expect("modulo result fits in usize");
Some(online[idx].idx())
}
pub(crate) fn snapshot(&self) -> view::Topology {
let mut view_tiers = Vec::with_capacity(self.pgproto_tiers.len());
add_to_view(&mut view_tiers, "pgproto", &self.pgproto_tiers);
for (connection_type, tiers) in &self.other_tiers {
add_to_view(&mut view_tiers, connection_type, tiers);
}
view_tiers.sort_unstable_by(|a, b| a.name.cmp(&b.name));
for tier in &mut view_tiers {
sort(&mut tier.replicasets);
}
view::Topology { tiers: view_tiers }
}
pub(crate) fn bucket_count(&self, tier: &str) -> Option<u64> {
self.pgproto_tiers
.iter()
.find(|t| t.name == tier)
.map(|t| t.bucket_count)
}
}
fn add_to_view(view_tiers: &mut Vec<view::Tier>, connection_type: &str, tiers: &[Tier]) {
for tier in tiers {
let view_tier = if let Some(t) = view_tiers.iter_mut().find(|t| t.name == tier.name) {
t
} else {
view_tiers.push(view::Tier {
name: tier.name.clone(),
bucket_count: tier.bucket_count,
replicasets: Vec::new(),
});
view_tiers.last_mut().expect("just pushed")
};
for (idx, rs) in tier.replicasets.iter().enumerate() {
let view_rs =
if let Some(r) = view_tier.replicasets.iter_mut().find(|r| r.uuid == rs.uuid) {
r
} else {
view_tier.replicasets.push(view::Replicaset {
uuid: rs.uuid,
name: rs.master.replicaset_name.clone(),
master: None,
replicas: Vec::new(),
ranges: tier
.bucket_map
.iter()
.filter(|s| s.replicaset_idx == idx)
.map(|s| s.start..=s.end)
.collect(),
});
view_tier.replicasets.last_mut().expect("just pushed")
};
if rs.master.is_master {
match &mut view_rs.master {
Some(m) => {
m.addresses
.insert(connection_type.to_string(), rs.master.addr.to_string());
}
None => view_rs.master = Some(instance(&rs.master, connection_type)),
}
}
for r in &rs.replicas {
if let Some(i) = view_rs.replicas.iter_mut().find(|i| i.name == r.name) {
i.addresses
.insert(connection_type.to_string(), r.addr.to_string());
} else {
view_rs.replicas.push(instance(r, connection_type));
}
}
}
}
}
impl Tier {
fn lookup(&self, bucket_id: u64) -> Option<&Replicaset> {
let i = self
.bucket_map
.partition_point(|s| s.start <= bucket_id)
.checked_sub(1)?;
let slot = &self.bucket_map[i];
if bucket_id <= slot.end {
self.replicasets.get(slot.replicaset_idx)
} else {
None
}
}
}
fn sort(replicasets: &mut [view::Replicaset]) {
replicasets.sort_unstable_by_key(|rs| rs.uuid);
for rs in replicasets {
rs.replicas.sort_unstable_by(|a, b| a.name.cmp(&b.name));
}
}
fn instance(inst: &Instance, connection_type: &str) -> view::Instance {
view::Instance {
name: inst.name.clone(),
addresses: HashMap::from([(connection_type.to_string(), inst.addr.to_string())]),
online: inst.is_online(),
is_master: inst.is_master,
raft_id: inst.raft_id,
uuid: inst.uuid,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::error::Error;
const RS1: Uuid = Uuid::from_u128(1);
const RS2: Uuid = Uuid::from_u128(2);
fn make_topology(name: &str, tier: &str, rs_uuid: Uuid) -> Topology {
let inst = Arc::new(Instance::new_for_test(
name,
"127.0.0.1:3301",
tier,
rs_uuid,
true,
3000,
));
let mut instances = HashMap::new();
instances.insert(inst.name.clone(), inst.clone());
let rs = Replicaset {
uuid: rs_uuid,
master: inst.clone(),
replicas: vec![],
};
let tier_obj = Tier {
name: tier.to_string(),
bucket_count: 3000,
replicasets: vec![rs],
bucket_map: vec![BucketSlot {
start: 1,
end: 3000,
replicaset_idx: 0,
}],
};
Topology {
instances,
pgproto_tiers: vec![tier_obj],
other_tiers: HashMap::new(),
}
}
#[test]
fn bucket_map_rejects_an_incomplete_map() {
let mut topo = Topology::new();
let mut incoming = make_topology("i1", "default", RS1);
incoming.pgproto_tiers[0].bucket_map = vec![
BucketSlot {
start: 1,
end: 1500,
replicaset_idx: 0,
},
BucketSlot {
start: 1501,
end: 3000,
replicaset_idx: 7,
},
];
topo.merge(incoming);
match topo.snapshot().tiers()[0].bucket_map() {
Err(Error::IncompleteBucketMap(m)) => {
assert_eq!(m.tier, "default");
assert_eq!((m.covered, m.expected), (1500, 3000));
}
other => panic!("a half-covered tier must not be returned: {other:?}"),
}
}
#[test]
fn bucket_map_rejects_a_tier_with_no_ranges() {
let mut topo = Topology::new();
let mut incoming = make_topology("i1", "default", RS1);
incoming.pgproto_tiers[0].bucket_map.clear();
topo.merge(incoming);
assert!(matches!(
topo.snapshot().tiers()[0].bucket_map(),
Err(Error::IncompleteBucketMap(ref m)) if m.covered == 0
));
}
#[test]
fn snapshot_ranges_agree_with_resolve() {
let mut topo = Topology::new();
let mut incoming = make_topology("i1", "default", RS1);
let i2 = Arc::new(Instance::new_for_test(
"i2",
"127.0.0.1:3302",
"default",
RS2,
true,
3000,
));
incoming.instances.insert(i2.name.clone(), i2.clone());
let tier = &mut incoming.pgproto_tiers[0];
tier.replicasets.push(Replicaset {
uuid: RS2,
master: i2,
replicas: vec![],
});
tier.bucket_map = vec![
BucketSlot {
start: 1,
end: 1000,
replicaset_idx: 0,
},
BucketSlot {
start: 1001,
end: 2000,
replicaset_idx: 1,
},
BucketSlot {
start: 2001,
end: 3000,
replicaset_idx: 0,
},
];
topo.merge(incoming);
let snapshot = topo.snapshot();
let tier = &snapshot.tiers()[0];
assert_eq!(tier.name(), "default");
let owned: Vec<_> = tier
.replicasets()
.iter()
.map(|rs| {
(
rs.uuid(),
rs.master().map(view::Instance::name),
rs.ranges().len(),
)
})
.collect();
assert_eq!(owned, [(RS1, Some("i1"), 2), (RS2, Some("i2"), 1)]);
let bucket_map = tier.bucket_map().expect("map covers the tier");
assert_eq!(bucket_map[&RS1], [1..=1000, 2001..=3000]);
for rs in tier.replicasets() {
let expected = topo.instances[rs.master().expect("named master").name()].idx();
for r in rs.ranges() {
for bucket in [*r.start(), *r.end()] {
assert_eq!(topo.resolve("default", bucket, true), Some(expected));
}
}
}
}
#[test]
fn snapshot_reports_no_master_until_one_is_named() {
let inst = Arc::new(Instance::new_for_test(
"i1",
"127.0.0.1:3301",
"default",
RS1,
false,
3000,
));
let mut incoming = make_topology("i1", "default", RS1);
incoming.instances.insert(inst.name.clone(), inst.clone());
let rs = &mut incoming.pgproto_tiers[0].replicasets[0];
rs.master = inst.clone();
rs.replicas = vec![inst];
let mut topo = Topology::new();
topo.merge(incoming);
assert_eq!(
topo.resolve("default", 1, true),
Some(topo.instances["i1"].idx())
);
let snapshot = topo.snapshot();
let rs = &snapshot.tiers()[0].replicasets()[0];
assert_eq!(rs.master(), None);
let replicas: Vec<_> = rs.replicas().iter().map(view::Instance::name).collect();
assert_eq!(replicas, ["i1"]);
}
#[test]
fn replicaset_lists_the_master_first() {
let mut topo = make_topology("m", "default", RS1);
topo.pgproto_tiers[0].replicasets[0].replicas = ["b", "a"]
.map(|name| {
Arc::new(Instance::new_for_test(
name,
"127.0.0.1:3302",
"default",
RS1,
false,
3000,
))
})
.into();
let snapshot = topo.snapshot();
let rs = &snapshot.tiers()[0].replicasets()[0];
let instances: Vec<_> = rs.instances().map(|i| (i.name(), i.is_master())).collect();
assert_eq!(instances, [("m", true), ("a", false), ("b", false)]);
}
#[test]
fn merge_preserves_pool_idx_in_tiers() {
NEXT_IDX.fetch_add(1, Ordering::Relaxed);
let mut topo = Topology::new();
topo.merge(make_topology("i1", "default", RS1));
let expected_idx = topo.instances["i1"].idx();
assert!(expected_idx > 0, "precondition: idx must be non-zero");
topo.merge(make_topology("i1", "default", RS1));
assert_eq!(
topo.resolve("default", 1, true),
Some(expected_idx),
"resolve must return the stable pool slot, not the default 0 from the fresh Arc"
);
}
#[test]
fn merge_keeps_one_object_per_instance() {
let mut topo = Topology::new();
topo.merge(make_topology("i1", "default", RS1));
topo.merge(make_topology("i1", "default", RS1));
let master = &topo.pgproto_tiers[0].replicasets[0].master;
assert!(Arc::ptr_eq(&topo.instances["i1"], master));
}
fn with_tier(tier: Tier) -> Topology {
Topology {
instances: HashMap::new(),
pgproto_tiers: Vec::new(),
other_tiers: HashMap::from([("iproto".to_string(), vec![tier])]),
}
}
fn tier_at(addr: &str) -> Tier {
let inst = Arc::new(Instance::new_for_test(
"i1", addr, "default", RS1, true, 3000,
));
Tier {
name: "default".to_string(),
bucket_count: 3000,
replicasets: vec![Replicaset {
uuid: RS1,
master: inst,
replicas: vec![],
}],
bucket_map: vec![BucketSlot {
start: 1,
end: 3000,
replicaset_idx: 0,
}],
}
}
#[test]
fn snapshot_merges_the_addresses_of_an_instance() {
let mut topo = with_tier(tier_at("127.0.0.1:3301"));
topo.pgproto_tiers = vec![tier_at("127.0.0.1:5432")];
let snapshot = topo.snapshot();
let rs = &snapshot.tiers()[0].replicasets()[0];
assert!(rs.replicas().is_empty());
let i1 = rs.master().expect("named master");
assert_eq!(i1.address("pgproto"), Some("127.0.0.1:5432"));
assert_eq!(i1.address("iproto"), Some("127.0.0.1:3301"));
}
#[test]
fn snapshot_keeps_an_instance_without_pgproto() {
let topo = with_tier(tier_at("127.0.0.1:3301"));
let snapshot = topo.snapshot();
let tier = &snapshot.tiers()[0];
let i1 = tier.replicasets()[0].master().expect("named master");
assert_eq!(i1.address("pgproto"), None);
assert_eq!(i1.address("iproto"), Some("127.0.0.1:3301"));
assert!(tier.bucket_map().is_ok());
}
#[test]
fn merge_replaces_other_tiers() {
let mut topo = with_tier(tier_at("127.0.0.1:3301"));
topo.merge(with_tier(tier_at("127.0.0.1:3302")));
let snapshot = topo.snapshot();
let i1 = snapshot.tiers()[0].replicasets()[0]
.master()
.expect("named master");
assert_eq!(i1.address("iproto"), Some("127.0.0.1:3302"));
}
#[test]
fn snapshot_is_independent_of_insertion_order() {
let inst = |name: &str, rs: Uuid, master: bool| {
Arc::new(Instance::new_for_test(
name,
"127.0.0.1:3301",
"default",
rs,
master,
3000,
))
};
let build = |reversed: bool| {
let mut sets = vec![
(RS1, 1, 1500, vec!["a", "b"]),
(RS2, 1501, 3000, vec!["c", "d"]),
];
if reversed {
sets.reverse();
}
let mut tier = Tier {
name: "default".to_string(),
bucket_count: 3000,
replicasets: Vec::new(),
bucket_map: Vec::new(),
};
for (idx, (uuid, start, end, mut names)) in sets.into_iter().enumerate() {
if reversed {
names.reverse();
}
tier.replicasets.push(Replicaset {
uuid,
master: inst(&format!("{uuid}-m"), uuid, true),
replicas: names.into_iter().map(|n| inst(n, uuid, false)).collect(),
});
tier.bucket_map.push(BucketSlot {
start,
end,
replicaset_idx: idx,
});
}
tier.bucket_map.sort_unstable_by_key(|s| s.start);
let mut topo = Topology::new();
topo.pgproto_tiers = vec![tier];
topo.snapshot()
};
assert_eq!(build(false), build(true));
}
}