#[cfg(test)]
use mockall::*;
use crate::cluster::topology::NodeDistance;
use crate::cluster::NodeInfo;
#[cfg_attr(test, automock)]
pub trait NodeDistanceEvaluator {
fn compute_distance(&self, node: &NodeInfo) -> Option<NodeDistance>;
}
#[derive(Default, Debug)]
pub struct AllLocalNodeDistanceEvaluator;
impl NodeDistanceEvaluator for AllLocalNodeDistanceEvaluator {
fn compute_distance(&self, _node: &NodeInfo) -> Option<NodeDistance> {
Some(NodeDistance::Local)
}
}
#[derive(Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Debug)]
pub struct TopologyAwareNodeDistanceEvaluator {
local_dc: String,
}
impl NodeDistanceEvaluator for TopologyAwareNodeDistanceEvaluator {
fn compute_distance(&self, node: &NodeInfo) -> Option<NodeDistance> {
Some(if node.datacenter == self.local_dc {
NodeDistance::Local
} else {
NodeDistance::Remote
})
}
}
impl TopologyAwareNodeDistanceEvaluator {
pub fn new(local_dc: String) -> Self {
TopologyAwareNodeDistanceEvaluator { local_dc }
}
}
#[cfg(test)]
mod tests {
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use uuid::Uuid;
use crate::cluster::topology::NodeDistance;
use crate::cluster::NodeInfo;
use crate::load_balancing::node_distance_evaluator::NodeDistanceEvaluator;
use crate::load_balancing::node_distance_evaluator::TopologyAwareNodeDistanceEvaluator;
#[test]
fn should_return_topology_aware_distance() {
let local_dc = "test";
let evaluator = TopologyAwareNodeDistanceEvaluator::new(local_dc.into());
assert_eq!(
evaluator
.compute_distance(&NodeInfo::new(
Uuid::new_v4(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080),
None,
"".into(),
Default::default(),
"".into(),
))
.unwrap(),
NodeDistance::Remote
);
assert_eq!(
evaluator
.compute_distance(&NodeInfo::new(
Uuid::new_v4(),
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080),
None,
local_dc.into(),
Default::default(),
"".into(),
))
.unwrap(),
NodeDistance::Local
);
}
}