use crate::hashkit::{random_slicing::RandomSlices, DynToken};
#[derive(Clone, Debug, Default)]
pub enum RackRing {
#[default]
Continuum,
RandomSlicing(RandomSlices),
}
#[derive(Clone, Debug)]
pub struct Continuum {
pub token: DynToken,
pub peer_idx: u32,
}
impl Continuum {
#[must_use]
pub fn new(token: DynToken, peer_idx: u32) -> Self {
Self { token, peer_idx }
}
}
#[derive(Clone, Debug)]
pub struct Rack {
name: String,
dc: String,
nserver_continuum: u32,
ncontinuum: u32,
continuums: Vec<Continuum>,
ring: RackRing,
}
impl Rack {
#[must_use]
pub fn new(name: String, dc: String) -> Self {
Self {
name,
dc,
nserver_continuum: 0,
ncontinuum: 0,
continuums: Vec::new(),
ring: RackRing::Continuum,
}
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn dc(&self) -> &str {
&self.dc
}
#[must_use]
pub fn continuums(&self) -> &[Continuum] {
&self.continuums
}
#[must_use]
pub fn nserver_continuum(&self) -> u32 {
self.nserver_continuum
}
#[must_use]
pub fn ncontinuum(&self) -> u32 {
self.ncontinuum
}
pub fn add_peer_tokens(&mut self, peer_idx: u32, tokens: &[DynToken]) {
for tok in tokens {
self.continuums.push(Continuum::new(tok.clone(), peer_idx));
self.ncontinuum = self.ncontinuum.saturating_add(1);
self.nserver_continuum = self.nserver_continuum.saturating_add(1);
}
}
pub fn sort_continuums(&mut self) {
self.continuums.sort_by(|a, b| a.token.cmp(&b.token));
}
pub fn clear_continuums(&mut self) {
self.continuums.clear();
self.ncontinuum = 0;
self.nserver_continuum = 0;
self.ring = RackRing::Continuum;
}
#[must_use]
pub fn ring(&self) -> &RackRing {
&self.ring
}
pub fn set_random_slices(&mut self, slices: RandomSlices) {
self.ring = RackRing::RandomSlicing(slices);
}
#[must_use]
pub fn is_random_slicing(&self) -> bool {
matches!(self.ring, RackRing::RandomSlicing(_))
}
#[must_use]
pub fn random_slices(&self) -> Option<&RandomSlices> {
match &self.ring {
RackRing::RandomSlicing(s) => Some(s),
RackRing::Continuum => None,
}
}
}
#[derive(Clone, Debug)]
pub struct Datacenter {
name: String,
racks: Vec<Rack>,
preselected_rack_for_replication: Option<usize>,
}
impl Datacenter {
#[must_use]
pub fn new(name: String) -> Self {
Self {
name,
racks: Vec::new(),
preselected_rack_for_replication: None,
}
}
#[must_use]
pub fn name(&self) -> &str {
&self.name
}
#[must_use]
pub fn racks(&self) -> &[Rack] {
&self.racks
}
pub fn racks_mut(&mut self) -> &mut [Rack] {
&mut self.racks
}
#[must_use]
pub fn rack(&self, name: &str) -> Option<&Rack> {
self.racks.iter().find(|r| r.name() == name)
}
pub fn rack_mut(&mut self, name: &str) -> Option<&mut Rack> {
self.racks.iter_mut().find(|r| r.name() == name)
}
#[must_use]
pub fn rack_idx(&self, name: &str) -> Option<usize> {
self.racks.iter().position(|r| r.name() == name)
}
pub fn upsert_rack(&mut self, name: String) -> &mut Rack {
if let Some(idx) = self.rack_idx(&name) {
return &mut self.racks[idx];
}
let dc = self.name.clone();
self.racks.push(Rack::new(name, dc));
let last = self.racks.len() - 1;
&mut self.racks[last]
}
pub fn sort_racks(&mut self) {
self.racks.sort_by(|a, b| a.name().cmp(b.name()));
}
#[must_use]
pub fn preselected_rack_idx(&self) -> Option<usize> {
self.preselected_rack_for_replication
}
#[must_use]
pub fn preselected_rack(&self) -> Option<&Rack> {
self.preselected_rack_for_replication
.and_then(|i| self.racks.get(i))
}
pub fn set_preselected_rack_idx(&mut self, idx: Option<usize>) {
self.preselected_rack_for_replication = idx;
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn upsert_is_idempotent() {
let mut dc = Datacenter::new("dc1".into());
dc.upsert_rack("r1".into());
dc.upsert_rack("r1".into());
assert_eq!(dc.racks().len(), 1);
}
#[test]
fn rack_continuum_sorts_by_token() {
let mut r = Rack::new("r".into(), "d".into());
r.add_peer_tokens(0, &[DynToken::from_u32(9)]);
r.add_peer_tokens(1, &[DynToken::from_u32(3)]);
r.add_peer_tokens(2, &[DynToken::from_u32(6)]);
r.sort_continuums();
let idxs: Vec<u32> = r.continuums().iter().map(|c| c.peer_idx).collect();
assert_eq!(idxs, vec![1, 2, 0]);
}
#[test]
fn rack_clear_resets_counters() {
let mut r = Rack::new("r".into(), "d".into());
r.add_peer_tokens(0, &[DynToken::from_u32(1)]);
r.clear_continuums();
assert_eq!(r.ncontinuum(), 0);
assert!(r.continuums().is_empty());
}
#[test]
fn sort_racks_alphabetical() {
let mut dc = Datacenter::new("dc1".into());
dc.upsert_rack("rb".into());
dc.upsert_rack("ra".into());
dc.sort_racks();
assert_eq!(dc.racks()[0].name(), "ra");
}
}