engine/distributed/

sharding.rs

//! Sharding Strategies for Distributed Vector Storage
//!
//! Supports multiple partitioning strategies:
//! - Consistent hashing for even distribution
//! - Range-based for locality-aware placement
//! - Custom strategies for specialized workloads

use serde::{Deserialize, Serialize};
use std::collections::hash_map::DefaultHasher;
use std::collections::{BTreeMap, HashMap};
use std::hash::{Hash, Hasher};

/// Configuration for sharding behavior
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardingConfig {
    /// Number of shards (partitions)
    pub num_shards: u32,
    /// Replication factor for each shard
    pub replication_factor: u32,
    /// Sharding strategy to use
    pub strategy: ShardingStrategy,
    /// Number of virtual nodes for consistent hashing
    pub virtual_nodes: u32,
}

impl Default for ShardingConfig {
    fn default() -> Self {
        Self {
            num_shards: 4,
            replication_factor: 2,
            strategy: ShardingStrategy::ConsistentHash,
            virtual_nodes: 150,
        }
    }
}

/// Strategy for distributing data across shards
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ShardingStrategy {
    /// Consistent hashing with virtual nodes
    ConsistentHash,
    /// Range-based partitioning
    Range,
    /// Simple modulo-based hashing
    Modulo,
}

/// Information about a single partition/shard
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PartitionInfo {
    /// Unique shard identifier
    pub shard_id: u32,
    /// Node IDs hosting this shard (primary + replicas)
    pub node_ids: Vec<String>,
    /// Primary node for writes
    pub primary_node: String,
    /// Whether this shard is healthy
    pub is_healthy: bool,
    /// Number of vectors in this shard
    pub vector_count: u64,
    /// Approximate memory usage in bytes
    pub memory_bytes: u64,
}

/// Assignment of a key to a shard
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ShardAssignment {
    /// The shard this key belongs to
    pub shard_id: u32,
    /// Node IDs that can serve this shard
    pub nodes: Vec<String>,
    /// Preferred node for this request
    pub preferred_node: String,
}

/// Consistent hash ring for shard assignment
#[derive(Debug, Clone)]
pub struct ConsistentHashRing {
    /// Ring of virtual node positions to shard IDs
    ring: BTreeMap<u64, u32>,
    /// Configuration
    config: ShardingConfig,
    /// Shard to nodes mapping
    shard_nodes: HashMap<u32, Vec<String>>,
}

impl ConsistentHashRing {
    /// Create a new consistent hash ring
    pub fn new(config: ShardingConfig) -> Self {
        let mut ring = BTreeMap::new();

        // Add virtual nodes for each shard
        for shard_id in 0..config.num_shards {
            for vnode in 0..config.virtual_nodes {
                let key = format!("shard-{}-vnode-{}", shard_id, vnode);
                let hash = Self::hash_key(&key);
                ring.insert(hash, shard_id);
            }
        }

        Self {
            ring,
            config,
            shard_nodes: HashMap::new(),
        }
    }

    /// Hash a key to a u64 position on the ring
    fn hash_key(key: &str) -> u64 {
        let mut hasher = DefaultHasher::new();
        key.hash(&mut hasher);
        hasher.finish()
    }

    /// Get the shard assignment for a vector ID
    pub fn get_shard(&self, vector_id: &str) -> ShardAssignment {
        let hash = Self::hash_key(vector_id);

        // Find the first node position >= hash
        let shard_id = self
            .ring
            .range(hash..)
            .next()
            .or_else(|| self.ring.iter().next())
            .map(|(_, &shard)| shard)
            .unwrap_or(0);

        let nodes = self
            .shard_nodes
            .get(&shard_id)
            .cloned()
            .unwrap_or_else(|| vec![format!("node-{}", shard_id)]);

        let preferred_node = nodes.first().cloned().unwrap_or_default();

        ShardAssignment {
            shard_id,
            nodes,
            preferred_node,
        }
    }

    /// Get shards for a batch of vector IDs
    pub fn get_shards_batch(&self, vector_ids: &[String]) -> HashMap<u32, Vec<String>> {
        let mut shard_vectors: HashMap<u32, Vec<String>> = HashMap::new();

        for id in vector_ids {
            let assignment = self.get_shard(id);
            shard_vectors
                .entry(assignment.shard_id)
                .or_default()
                .push(id.clone());
        }

        shard_vectors
    }

    /// Register nodes for a shard
    pub fn register_shard_nodes(&mut self, shard_id: u32, node_ids: Vec<String>) {
        self.shard_nodes.insert(shard_id, node_ids);
    }

    /// Get all shard IDs
    pub fn get_all_shards(&self) -> Vec<u32> {
        (0..self.config.num_shards).collect()
    }

    /// Get partition info for all shards
    pub fn get_partition_info(&self) -> Vec<PartitionInfo> {
        (0..self.config.num_shards)
            .map(|shard_id| {
                let nodes = self
                    .shard_nodes
                    .get(&shard_id)
                    .cloned()
                    .unwrap_or_else(|| vec![format!("node-{}", shard_id)]);
                let primary = nodes.first().cloned().unwrap_or_default();

                PartitionInfo {
                    shard_id,
                    node_ids: nodes,
                    primary_node: primary,
                    is_healthy: true,
                    vector_count: 0,
                    memory_bytes: 0,
                }
            })
            .collect()
    }

    /// Rebalance when nodes change
    pub fn rebalance(&mut self, new_node_count: u32) {
        // Redistribute shards across new node count
        for shard_id in 0..self.config.num_shards {
            let mut nodes = Vec::new();
            for replica in 0..self.config.replication_factor.min(new_node_count) {
                let node_idx = (shard_id + replica) % new_node_count;
                nodes.push(format!("node-{}", node_idx));
            }
            self.shard_nodes.insert(shard_id, nodes);
        }
    }
}

/// Range-based sharding for ordered data
#[derive(Debug, Clone)]
pub struct RangeSharder {
    /// Boundaries for each shard (exclusive upper bounds)
    boundaries: Vec<u64>,
    /// Configuration
    config: ShardingConfig,
    /// Shard to nodes mapping
    shard_nodes: HashMap<u32, Vec<String>>,
}

impl RangeSharder {
    /// Create a new range sharder with even distribution
    pub fn new(config: ShardingConfig) -> Self {
        let step = u64::MAX / config.num_shards as u64;
        let boundaries: Vec<u64> = (1..config.num_shards).map(|i| step * i as u64).collect();

        Self {
            boundaries,
            config,
            shard_nodes: HashMap::new(),
        }
    }

    /// Get shard for a vector ID using range partitioning
    pub fn get_shard(&self, vector_id: &str) -> ShardAssignment {
        let hash = {
            let mut hasher = DefaultHasher::new();
            vector_id.hash(&mut hasher);
            hasher.finish()
        };

        // Find which range this hash falls into
        let shard_id = self
            .boundaries
            .iter()
            .position(|&b| hash < b)
            .unwrap_or(self.config.num_shards as usize - 1) as u32;

        let nodes = self
            .shard_nodes
            .get(&shard_id)
            .cloned()
            .unwrap_or_else(|| vec![format!("node-{}", shard_id)]);

        let preferred_node = nodes.first().cloned().unwrap_or_default();

        ShardAssignment {
            shard_id,
            nodes,
            preferred_node,
        }
    }

    /// Register nodes for a shard
    pub fn register_shard_nodes(&mut self, shard_id: u32, node_ids: Vec<String>) {
        self.shard_nodes.insert(shard_id, node_ids);
    }
}

/// Unified shard manager supporting multiple strategies
pub struct ShardManager {
    config: ShardingConfig,
    consistent_ring: Option<ConsistentHashRing>,
    range_sharder: Option<RangeSharder>,
}

impl ShardManager {
    /// Create a new shard manager
    pub fn new(config: ShardingConfig) -> Self {
        let (consistent_ring, range_sharder) = match config.strategy {
            ShardingStrategy::ConsistentHash | ShardingStrategy::Modulo => {
                (Some(ConsistentHashRing::new(config.clone())), None)
            }
            ShardingStrategy::Range => (None, Some(RangeSharder::new(config.clone()))),
        };

        Self {
            config,
            consistent_ring,
            range_sharder,
        }
    }

    /// Get shard assignment for a vector ID
    pub fn get_shard(&self, vector_id: &str) -> ShardAssignment {
        match self.config.strategy {
            ShardingStrategy::ConsistentHash | ShardingStrategy::Modulo => {
                match self.consistent_ring.as_ref() {
                    Some(ring) => ring.get_shard(vector_id),
                    None => {
                        tracing::error!("consistent_ring not initialized for ConsistentHash/Modulo strategy — falling back to shard 0");
                        ShardAssignment {
                            shard_id: 0,
                            nodes: vec![],
                            preferred_node: String::new(),
                        }
                    }
                }
            }
            ShardingStrategy::Range => match self.range_sharder.as_ref() {
                Some(sharder) => sharder.get_shard(vector_id),
                None => {
                    tracing::error!("range_sharder not initialized for Range strategy — falling back to shard 0");
                    ShardAssignment {
                        shard_id: 0,
                        nodes: vec![],
                        preferred_node: String::new(),
                    }
                }
            },
        }
    }

    /// Get shards for a batch of vector IDs
    pub fn get_shards_batch(&self, vector_ids: &[String]) -> HashMap<u32, Vec<String>> {
        let mut shard_vectors: HashMap<u32, Vec<String>> = HashMap::new();

        for id in vector_ids {
            let assignment = self.get_shard(id);
            shard_vectors
                .entry(assignment.shard_id)
                .or_default()
                .push(id.clone());
        }

        shard_vectors
    }

    /// Get all shard IDs for scatter queries
    pub fn get_all_shards(&self) -> Vec<u32> {
        (0..self.config.num_shards).collect()
    }

    /// Register nodes for a shard
    pub fn register_shard_nodes(&mut self, shard_id: u32, node_ids: Vec<String>) {
        if let Some(ref mut ring) = self.consistent_ring {
            ring.register_shard_nodes(shard_id, node_ids);
        } else if let Some(ref mut sharder) = self.range_sharder {
            sharder.register_shard_nodes(shard_id, node_ids);
        }
    }

    /// Get partition information for all shards
    pub fn get_partition_info(&self) -> Vec<PartitionInfo> {
        if let Some(ref ring) = self.consistent_ring {
            ring.get_partition_info()
        } else {
            (0..self.config.num_shards)
                .map(|shard_id| PartitionInfo {
                    shard_id,
                    node_ids: vec![format!("node-{}", shard_id)],
                    primary_node: format!("node-{}", shard_id),
                    is_healthy: true,
                    vector_count: 0,
                    memory_bytes: 0,
                })
                .collect()
        }
    }

    /// Rebalance shards across nodes
    pub fn rebalance(&mut self, node_count: u32) {
        if let Some(ref mut ring) = self.consistent_ring {
            ring.rebalance(node_count);
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_consistent_hash_ring() {
        let config = ShardingConfig {
            num_shards: 4,
            replication_factor: 2,
            strategy: ShardingStrategy::ConsistentHash,
            virtual_nodes: 100,
        };

        let ring = ConsistentHashRing::new(config);

        // Test that same key always maps to same shard
        let assignment1 = ring.get_shard("vector-123");
        let assignment2 = ring.get_shard("vector-123");
        assert_eq!(assignment1.shard_id, assignment2.shard_id);

        // Test that shards are in valid range
        for i in 0..100 {
            let assignment = ring.get_shard(&format!("test-{}", i));
            assert!(assignment.shard_id < 4);
        }
    }

    #[test]
    fn test_consistent_hash_distribution() {
        let config = ShardingConfig {
            num_shards: 4,
            replication_factor: 2,
            strategy: ShardingStrategy::ConsistentHash,
            virtual_nodes: 150,
        };

        let ring = ConsistentHashRing::new(config);

        // Count distribution across shards
        let mut counts = [0u32; 4];
        for i in 0..1000 {
            let assignment = ring.get_shard(&format!("vector-{}", i));
            counts[assignment.shard_id as usize] += 1;
        }

        // Check reasonably even distribution (within 50% of average)
        let avg = 250.0;
        for count in counts {
            assert!(count as f64 > avg * 0.5);
            assert!((count as f64) < avg * 1.5);
        }
    }

    #[test]
    fn test_batch_sharding() {
        let config = ShardingConfig::default();
        let ring = ConsistentHashRing::new(config);

        let ids: Vec<String> = (0..100).map(|i| format!("vec-{}", i)).collect();
        let shard_batches = ring.get_shards_batch(&ids);

        // All IDs should be distributed
        let total: usize = shard_batches.values().map(|v| v.len()).sum();
        assert_eq!(total, 100);
    }

    #[test]
    fn test_range_sharder() {
        let config = ShardingConfig {
            num_shards: 4,
            replication_factor: 1,
            strategy: ShardingStrategy::Range,
            virtual_nodes: 0, // Not used for range
        };

        let sharder = RangeSharder::new(config);

        // Test deterministic assignment
        let a1 = sharder.get_shard("test-key");
        let a2 = sharder.get_shard("test-key");
        assert_eq!(a1.shard_id, a2.shard_id);

        // Test shard range
        for i in 0..100 {
            let assignment = sharder.get_shard(&format!("key-{}", i));
            assert!(assignment.shard_id < 4);
        }
    }

    #[test]
    fn test_shard_manager() {
        let config = ShardingConfig::default();
        let mut manager = ShardManager::new(config);

        // Register nodes
        manager.register_shard_nodes(0, vec!["node-a".to_string(), "node-b".to_string()]);

        // Test sharding
        let assignment = manager.get_shard("my-vector");
        assert!(assignment.shard_id < 4);

        // Test all shards
        let shards = manager.get_all_shards();
        assert_eq!(shards.len(), 4);

        // Test partition info
        let partitions = manager.get_partition_info();
        assert_eq!(partitions.len(), 4);
    }

    #[test]
    fn test_rebalance() {
        let config = ShardingConfig {
            num_shards: 4,
            replication_factor: 2,
            ..Default::default()
        };

        let mut ring = ConsistentHashRing::new(config);
        ring.rebalance(3);

        // After rebalance, each shard should have nodes assigned
        let partitions = ring.get_partition_info();
        for partition in partitions {
            assert!(!partition.node_ids.is_empty());
            assert!(partition.node_ids.len() <= 2); // replication_factor
        }
    }
}