rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation
use crate::cache_redis::{RedisCacheError, RedisCacheResult, RedisNodeConfig};

/// Public shard metadata used by the Redis rendezvous picker.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisShardChoice {
    /// Stable node name.
    pub name: String,
    /// Relative node weight.
    pub weight: u32,
}

impl From<&RedisNodeConfig> for RedisShardChoice {
    fn from(value: &RedisNodeConfig) -> Self {
        Self {
            name: value.name.clone(),
            weight: value.weight,
        }
    }
}

/// Weighted rendezvous hashing picker for Redis application-level sharding.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RedisShardPicker {
    nodes: Vec<RedisShardChoice>,
}

impl RedisShardPicker {
    /// Creates a picker from validated shard metadata.
    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 })
    }

    /// Returns the selected shard name for a cache key.
    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}");
    }
}