use crate::cache_redis::{RedisCacheError, RedisCacheResult, RedisNodeConfig};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisShardChoice {
pub name: String,
pub weight: u32,
}
impl From<&RedisNodeConfig> for RedisShardChoice {
fn from(value: &RedisNodeConfig) -> Self {
Self {
name: value.name.clone(),
weight: value.weight,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisShardPicker {
nodes: Vec<RedisShardChoice>,
}
impl RedisShardPicker {
pub fn new(nodes: Vec<RedisShardChoice>) -> RedisCacheResult<Self> {
if nodes.is_empty() {
return Err(RedisCacheError::InvalidConfig(
"at least one redis shard is required".to_string(),
));
}
if let Some(node) = nodes.iter().find(|node| node.weight == 0) {
return Err(RedisCacheError::InvalidConfig(format!(
"redis shard `{}` weight must be greater than zero",
node.name
)));
}
Ok(Self { nodes })
}
pub fn select_name(&self, key: &str) -> RedisCacheResult<&str> {
self.nodes
.iter()
.max_by_key(|node| self.weighted_score(key, node))
.map(|node| node.name.as_str())
.ok_or_else(|| {
RedisCacheError::InvalidConfig("at least one redis shard is required".to_string())
})
}
fn weighted_score(&self, key: &str, node: &RedisShardChoice) -> u64 {
let mut best = 0;
for replica in 0..node.weight {
best = best.max(fnv1a64(key, &node.name, replica));
}
best
}
}
fn fnv1a64(key: &str, node: &str, replica: u32) -> u64 {
const OFFSET: u64 = 0xcbf29ce484222325;
const PRIME: u64 = 0x100000001b3;
let mut hash = OFFSET;
for byte in key
.as_bytes()
.iter()
.chain([b'|'].iter())
.chain(node.as_bytes())
.chain([b'|'].iter())
.chain(replica.to_le_bytes().iter())
{
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(PRIME);
}
hash
}
#[cfg(test)]
mod tests {
use crate::cache_redis::{RedisShardChoice, RedisShardPicker};
#[test]
fn shard_picker_selects_same_node_for_same_key() {
let picker = RedisShardPicker::new(vec![
RedisShardChoice {
name: "a".to_string(),
weight: 1,
},
RedisShardChoice {
name: "b".to_string(),
weight: 1,
},
])
.expect("picker");
assert_eq!(
picker.select_name("user:1").expect("first"),
picker.select_name("user:1").expect("second")
);
}
#[test]
fn shard_picker_respects_weight_over_many_keys() {
let picker = RedisShardPicker::new(vec![
RedisShardChoice {
name: "light".to_string(),
weight: 1,
},
RedisShardChoice {
name: "heavy".to_string(),
weight: 8,
},
])
.expect("picker");
let mut heavy = 0;
let mut light = 0;
for index in 0..200 {
match picker.select_name(&format!("key:{index}")).expect("node") {
"heavy" => heavy += 1,
"light" => light += 1,
other => panic!("unexpected shard {other}"),
}
}
assert!(heavy > light, "heavy={heavy} light={light}");
}
}