use std::collections::BTreeMap;
use serde::{Deserialize, Serialize};
pub type Region = String;
pub const FAR: u32 = u32::MAX / 2;
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default, deny_unknown_fields)]
pub struct RegionPreference {
#[serde(skip_serializing_if = "Option::is_none")]
pub prefer: Option<Region>,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub allow: Vec<Region>,
}
impl RegionPreference {
pub fn allows(&self, region: &str) -> bool {
self.allow.is_empty() || self.allow.iter().any(|r| r == region)
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct RegionMap {
edges: BTreeMap<String, u32>,
}
impl RegionMap {
pub fn from_edges(edges: impl IntoIterator<Item = (Region, Region, u32)>) -> Self {
let mut map = BTreeMap::new();
for (a, b, d) in edges {
map.insert(Self::edge_key(&a, &b), d);
map.insert(Self::edge_key(&b, &a), d);
}
Self { edges: map }
}
pub fn distance(&self, a: &str, b: &str) -> u32 {
if a == b {
return 0;
}
self.edges
.get(&Self::edge_key(a, b))
.or_else(|| self.edges.get(&Self::edge_key(b, a)))
.copied()
.unwrap_or(FAR)
}
pub fn is_empty(&self) -> bool {
self.edges.is_empty()
}
fn edge_key(a: &str, b: &str) -> String {
format!("{a}\u{1}{b}")
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RegionCandidate {
pub region: Option<Region>,
pub healthy: bool,
}
pub fn rank_by_nearest(
candidates: &[RegionCandidate],
client: Option<&str>,
map: &RegionMap,
) -> Vec<usize> {
let mut order: Vec<usize> = (0..candidates.len()).collect();
order.sort_by_key(|&i| {
let c = &candidates[i];
let health_rank = if c.healthy { 0 } else { 1 };
let dist = match (client, &c.region) {
(Some(client), Some(region)) => map.distance(client, region),
_ => FAR / 2,
};
(health_rank, dist, i)
});
order
}
#[cfg(test)]
mod tests {
use super::*;
fn cand(region: Option<&str>, healthy: bool) -> RegionCandidate {
RegionCandidate {
region: region.map(String::from),
healthy,
}
}
#[test]
fn distance_is_symmetric_with_defaults() {
let map = RegionMap::from_edges([
("us-east".into(), "us-west".into(), 1),
("us-east".into(), "eu-west".into(), 3),
]);
assert_eq!(map.distance("us-east", "us-east"), 0);
assert_eq!(map.distance("us-east", "us-west"), 1);
assert_eq!(map.distance("us-west", "us-east"), 1); assert_eq!(map.distance("us-east", "eu-west"), 3);
assert_eq!(map.distance("us-east", "ap-south"), FAR); }
#[test]
fn nearest_prefers_same_region_then_by_distance() {
let candidates = [
cand(Some("eu-west"), true),
cand(Some("us-east"), true),
cand(Some("us-west"), true),
];
let map = RegionMap::from_edges([
("us-east".into(), "us-west".into(), 1),
("us-east".into(), "eu-west".into(), 3),
]);
assert_eq!(
rank_by_nearest(&candidates, Some("us-east"), &map),
vec![1, 2, 0]
);
}
#[test]
fn unhealthy_candidates_sort_last_even_if_nearer() {
let candidates = [
cand(Some("us-east"), false), cand(Some("eu-west"), true), ];
let map = RegionMap::from_edges([("us-east".into(), "eu-west".into(), 5)]);
let order = rank_by_nearest(&candidates, Some("us-east"), &map);
assert_eq!(order, vec![1, 0]); }
#[test]
fn no_client_region_orders_by_health_then_original() {
let candidates = [
cand(Some("eu-west"), false),
cand(Some("us-east"), true),
cand(None, true),
];
let map = RegionMap::default();
assert_eq!(rank_by_nearest(&candidates, None, &map), vec![1, 2, 0]);
}
#[test]
fn region_preference_allow_list() {
let pref = RegionPreference {
prefer: Some("us-east".into()),
allow: vec!["us-east".into(), "us-west".into()],
};
assert!(pref.allows("us-east"));
assert!(!pref.allows("eu-west"));
assert!(RegionPreference::default().allows("anywhere"));
}
}