Merge commit 'd803bfe2b1fe7f5e219e50ac20d6801a0a58ac75' as 'vendor/ruvector'
This commit is contained in:
397
vendor/ruvector/tests/integration/distributed/sharding_tests.rs
vendored
Normal file
397
vendor/ruvector/tests/integration/distributed/sharding_tests.rs
vendored
Normal file
@@ -0,0 +1,397 @@
|
||||
//! Auto-Sharding Tests
|
||||
//!
|
||||
//! Tests for:
|
||||
//! - Consistent hashing for shard distribution
|
||||
//! - Dynamic shard rebalancing
|
||||
//! - Cross-shard queries
|
||||
//! - Load balancing
|
||||
|
||||
use ruvector_cluster::{
|
||||
ConsistentHashRing, ShardRouter, ClusterManager, ClusterConfig,
|
||||
ClusterNode, ShardInfo, ShardStatus, NodeStatus,
|
||||
discovery::StaticDiscovery,
|
||||
shard::{ShardMigration, LoadBalancer, LoadStats},
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
use std::net::{SocketAddr, IpAddr, Ipv4Addr};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
/// Test consistent hash ring creation
|
||||
#[tokio::test]
|
||||
async fn test_consistent_hash_ring_creation() {
|
||||
let ring = ConsistentHashRing::new(3);
|
||||
|
||||
assert_eq!(ring.node_count(), 0);
|
||||
assert!(ring.list_nodes().is_empty());
|
||||
}
|
||||
|
||||
/// Test adding nodes to hash ring
|
||||
#[tokio::test]
|
||||
async fn test_hash_ring_node_addition() {
|
||||
let mut ring = ConsistentHashRing::new(3);
|
||||
|
||||
ring.add_node("node1".to_string());
|
||||
ring.add_node("node2".to_string());
|
||||
ring.add_node("node3".to_string());
|
||||
|
||||
assert_eq!(ring.node_count(), 3);
|
||||
|
||||
let nodes = ring.list_nodes();
|
||||
assert!(nodes.contains(&"node1".to_string()));
|
||||
assert!(nodes.contains(&"node2".to_string()));
|
||||
assert!(nodes.contains(&"node3".to_string()));
|
||||
}
|
||||
|
||||
/// Test node removal from hash ring
|
||||
#[tokio::test]
|
||||
async fn test_hash_ring_node_removal() {
|
||||
let mut ring = ConsistentHashRing::new(3);
|
||||
|
||||
ring.add_node("node1".to_string());
|
||||
ring.add_node("node2".to_string());
|
||||
ring.add_node("node3".to_string());
|
||||
|
||||
assert_eq!(ring.node_count(), 3);
|
||||
|
||||
ring.remove_node("node2");
|
||||
|
||||
assert_eq!(ring.node_count(), 2);
|
||||
assert!(!ring.list_nodes().contains(&"node2".to_string()));
|
||||
}
|
||||
|
||||
/// Test key distribution across nodes
|
||||
#[tokio::test]
|
||||
async fn test_key_distribution() {
|
||||
let mut ring = ConsistentHashRing::new(3);
|
||||
|
||||
ring.add_node("node1".to_string());
|
||||
ring.add_node("node2".to_string());
|
||||
ring.add_node("node3".to_string());
|
||||
|
||||
let mut distribution: HashMap<String, usize> = HashMap::new();
|
||||
|
||||
// Test distribution across many keys
|
||||
for i in 0..10000 {
|
||||
let key = format!("key-{}", i);
|
||||
if let Some(node) = ring.get_primary_node(&key) {
|
||||
*distribution.entry(node).or_insert(0) += 1;
|
||||
}
|
||||
}
|
||||
|
||||
println!("Key distribution across nodes:");
|
||||
for (node, count) in &distribution {
|
||||
let percentage = (*count as f64 / 10000.0) * 100.0;
|
||||
println!(" {}: {} ({:.1}%)", node, count, percentage);
|
||||
}
|
||||
|
||||
// Each node should get roughly 1/3 (within reasonable tolerance)
|
||||
for count in distribution.values() {
|
||||
let ratio = *count as f64 / 10000.0;
|
||||
assert!(ratio > 0.2 && ratio < 0.5, "Uneven distribution: {:.3}", ratio);
|
||||
}
|
||||
}
|
||||
|
||||
/// Test replication factor compliance
|
||||
#[tokio::test]
|
||||
async fn test_replication_factor() {
|
||||
let mut ring = ConsistentHashRing::new(3);
|
||||
|
||||
ring.add_node("node1".to_string());
|
||||
ring.add_node("node2".to_string());
|
||||
ring.add_node("node3".to_string());
|
||||
ring.add_node("node4".to_string());
|
||||
ring.add_node("node5".to_string());
|
||||
|
||||
// Request 3 nodes for replication
|
||||
let nodes = ring.get_nodes("test-key", 3);
|
||||
|
||||
assert_eq!(nodes.len(), 3);
|
||||
|
||||
// All nodes should be unique
|
||||
let unique: std::collections::HashSet<_> = nodes.iter().collect();
|
||||
assert_eq!(unique.len(), 3);
|
||||
}
|
||||
|
||||
/// Test shard router creation
|
||||
#[tokio::test]
|
||||
async fn test_shard_router() {
|
||||
let router = ShardRouter::new(64);
|
||||
|
||||
let shard1 = router.get_shard("key1");
|
||||
let shard2 = router.get_shard("key2");
|
||||
|
||||
assert!(shard1 < 64);
|
||||
assert!(shard2 < 64);
|
||||
|
||||
// Same key should always map to same shard
|
||||
let shard1_again = router.get_shard("key1");
|
||||
assert_eq!(shard1, shard1_again);
|
||||
}
|
||||
|
||||
/// Test jump consistent hash distribution
|
||||
#[tokio::test]
|
||||
async fn test_jump_consistent_hash() {
|
||||
let router = ShardRouter::new(16);
|
||||
|
||||
let mut distribution: HashMap<u32, usize> = HashMap::new();
|
||||
|
||||
for i in 0..10000 {
|
||||
let key = format!("test-key-{}", i);
|
||||
let shard = router.get_shard(&key);
|
||||
*distribution.entry(shard).or_insert(0) += 1;
|
||||
}
|
||||
|
||||
println!("Shard distribution:");
|
||||
let mut total = 0;
|
||||
for shard in 0..16 {
|
||||
let count = distribution.get(&shard).copied().unwrap_or(0);
|
||||
total += count;
|
||||
println!(" Shard {}: {}", shard, count);
|
||||
}
|
||||
|
||||
assert_eq!(total, 10000);
|
||||
|
||||
// Check for reasonably even distribution
|
||||
let expected = 10000 / 16;
|
||||
for count in distribution.values() {
|
||||
let deviation = (*count as i32 - expected as i32).abs() as f64 / expected as f64;
|
||||
assert!(deviation < 0.5, "Shard distribution too uneven");
|
||||
}
|
||||
}
|
||||
|
||||
/// Test shard router caching
|
||||
#[tokio::test]
|
||||
async fn test_shard_router_caching() {
|
||||
let router = ShardRouter::new(64);
|
||||
|
||||
// First access
|
||||
let _ = router.get_shard("cached-key");
|
||||
|
||||
let stats = router.cache_stats();
|
||||
assert_eq!(stats.entries, 1);
|
||||
|
||||
// Second access (should hit cache)
|
||||
let _ = router.get_shard("cached-key");
|
||||
|
||||
// Add more keys
|
||||
for i in 0..100 {
|
||||
router.get_shard(&format!("key-{}", i));
|
||||
}
|
||||
|
||||
let stats = router.cache_stats();
|
||||
assert_eq!(stats.entries, 101); // 1 original + 100 new
|
||||
}
|
||||
|
||||
/// Test cache clearing
|
||||
#[tokio::test]
|
||||
async fn test_cache_clearing() {
|
||||
let router = ShardRouter::new(32);
|
||||
|
||||
for i in 0..50 {
|
||||
router.get_shard(&format!("key-{}", i));
|
||||
}
|
||||
|
||||
assert_eq!(router.cache_stats().entries, 50);
|
||||
|
||||
router.clear_cache();
|
||||
|
||||
assert_eq!(router.cache_stats().entries, 0);
|
||||
}
|
||||
|
||||
/// Test cluster manager creation
|
||||
#[tokio::test]
|
||||
async fn test_cluster_manager_creation() {
|
||||
let config = ClusterConfig::default();
|
||||
let discovery = Box::new(StaticDiscovery::new(vec![]));
|
||||
|
||||
let manager = ClusterManager::new(config, "test-node".to_string(), discovery);
|
||||
|
||||
assert!(manager.is_ok());
|
||||
}
|
||||
|
||||
/// Test cluster node management
|
||||
#[tokio::test]
|
||||
async fn test_cluster_node_management() {
|
||||
let config = ClusterConfig {
|
||||
shard_count: 8,
|
||||
replication_factor: 2,
|
||||
..Default::default()
|
||||
};
|
||||
let discovery = Box::new(StaticDiscovery::new(vec![]));
|
||||
let manager = ClusterManager::new(config, "coordinator".to_string(), discovery).unwrap();
|
||||
|
||||
// Add nodes
|
||||
for i in 0..3 {
|
||||
let node = ClusterNode::new(
|
||||
format!("node{}", i),
|
||||
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9000 + i as u16),
|
||||
);
|
||||
manager.add_node(node).await.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(manager.list_nodes().len(), 3);
|
||||
|
||||
// Remove a node
|
||||
manager.remove_node("node1").await.unwrap();
|
||||
assert_eq!(manager.list_nodes().len(), 2);
|
||||
}
|
||||
|
||||
/// Test shard assignment
|
||||
#[tokio::test]
|
||||
async fn test_shard_assignment() {
|
||||
let config = ClusterConfig {
|
||||
shard_count: 4,
|
||||
replication_factor: 2,
|
||||
..Default::default()
|
||||
};
|
||||
let discovery = Box::new(StaticDiscovery::new(vec![]));
|
||||
let manager = ClusterManager::new(config, "coordinator".to_string(), discovery).unwrap();
|
||||
|
||||
// Add nodes
|
||||
for i in 0..3 {
|
||||
let node = ClusterNode::new(
|
||||
format!("node{}", i),
|
||||
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9000 + i as u16),
|
||||
);
|
||||
manager.add_node(node).await.unwrap();
|
||||
}
|
||||
|
||||
// Assign a shard
|
||||
let shard = manager.assign_shard(0).unwrap();
|
||||
|
||||
assert_eq!(shard.shard_id, 0);
|
||||
assert!(!shard.primary_node.is_empty());
|
||||
assert_eq!(shard.status, ShardStatus::Active);
|
||||
}
|
||||
|
||||
/// Test shard migration
|
||||
#[tokio::test]
|
||||
async fn test_shard_migration() {
|
||||
let mut migration = ShardMigration::new(0, 1, 1000);
|
||||
|
||||
assert!(!migration.is_complete());
|
||||
assert_eq!(migration.progress, 0.0);
|
||||
|
||||
// Simulate partial migration
|
||||
migration.update_progress(500);
|
||||
assert_eq!(migration.progress, 0.5);
|
||||
assert!(!migration.is_complete());
|
||||
|
||||
// Complete migration
|
||||
migration.update_progress(1000);
|
||||
assert_eq!(migration.progress, 1.0);
|
||||
assert!(migration.is_complete());
|
||||
}
|
||||
|
||||
/// Test load balancer
|
||||
#[tokio::test]
|
||||
async fn test_load_balancer() {
|
||||
let balancer = LoadBalancer::new();
|
||||
|
||||
// Update loads for shards
|
||||
balancer.update_load(0, 0.3);
|
||||
balancer.update_load(1, 0.8);
|
||||
balancer.update_load(2, 0.5);
|
||||
balancer.update_load(3, 0.2);
|
||||
|
||||
// Get loads
|
||||
assert_eq!(balancer.get_load(0), 0.3);
|
||||
assert_eq!(balancer.get_load(1), 0.8);
|
||||
|
||||
// Get least loaded shard
|
||||
let least_loaded = balancer.get_least_loaded_shard(&[0, 1, 2, 3]);
|
||||
assert_eq!(least_loaded, Some(3));
|
||||
|
||||
// Get statistics
|
||||
let stats = balancer.get_stats();
|
||||
assert_eq!(stats.shard_count, 4);
|
||||
assert!(stats.avg_load > 0.0);
|
||||
assert!(stats.max_load == 0.8);
|
||||
assert!(stats.min_load == 0.2);
|
||||
}
|
||||
|
||||
/// Test cluster statistics
|
||||
#[tokio::test]
|
||||
async fn test_cluster_statistics() {
|
||||
let config = ClusterConfig {
|
||||
shard_count: 4,
|
||||
replication_factor: 2,
|
||||
..Default::default()
|
||||
};
|
||||
let discovery = Box::new(StaticDiscovery::new(vec![]));
|
||||
let manager = ClusterManager::new(config, "coordinator".to_string(), discovery).unwrap();
|
||||
|
||||
// Add nodes
|
||||
for i in 0..3 {
|
||||
let node = ClusterNode::new(
|
||||
format!("node{}", i),
|
||||
SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 9000 + i as u16),
|
||||
);
|
||||
manager.add_node(node).await.unwrap();
|
||||
}
|
||||
|
||||
let stats = manager.get_stats();
|
||||
|
||||
assert_eq!(stats.total_nodes, 3);
|
||||
assert_eq!(stats.healthy_nodes, 3);
|
||||
}
|
||||
|
||||
/// Performance test for shard routing
|
||||
#[tokio::test]
|
||||
async fn test_shard_routing_performance() {
|
||||
let router = ShardRouter::new(256);
|
||||
|
||||
let start = Instant::now();
|
||||
let iterations = 100000;
|
||||
|
||||
for i in 0..iterations {
|
||||
let key = format!("perf-key-{}", i);
|
||||
let _ = router.get_shard(&key);
|
||||
}
|
||||
|
||||
let elapsed = start.elapsed();
|
||||
let ops_per_sec = iterations as f64 / elapsed.as_secs_f64();
|
||||
|
||||
println!("Shard routing: {:.0} ops/sec", ops_per_sec);
|
||||
|
||||
// Should be able to route millions of keys per second
|
||||
assert!(ops_per_sec > 100000.0, "Routing too slow: {:.0} ops/sec", ops_per_sec);
|
||||
}
|
||||
|
||||
/// Test key stability after node addition
|
||||
#[tokio::test]
|
||||
async fn test_key_stability_on_node_addition() {
|
||||
let mut ring = ConsistentHashRing::new(3);
|
||||
|
||||
ring.add_node("node1".to_string());
|
||||
ring.add_node("node2".to_string());
|
||||
ring.add_node("node3".to_string());
|
||||
|
||||
// Record initial assignments
|
||||
let mut initial_assignments: HashMap<String, String> = HashMap::new();
|
||||
for i in 0..1000 {
|
||||
let key = format!("stable-key-{}", i);
|
||||
if let Some(node) = ring.get_primary_node(&key) {
|
||||
initial_assignments.insert(key, node);
|
||||
}
|
||||
}
|
||||
|
||||
// Add a new node
|
||||
ring.add_node("node4".to_string());
|
||||
|
||||
// Check how many keys changed
|
||||
let mut changes = 0;
|
||||
for (key, original_node) in &initial_assignments {
|
||||
if let Some(new_node) = ring.get_primary_node(key) {
|
||||
if new_node != *original_node {
|
||||
changes += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let change_ratio = changes as f64 / initial_assignments.len() as f64;
|
||||
println!("Keys reassigned after adding node: {} ({:.1}%)", changes, change_ratio * 100.0);
|
||||
|
||||
// With consistent hashing, only ~25% of keys should move (1/4 for 3->4 nodes)
|
||||
assert!(change_ratio < 0.4, "Too many keys reassigned: {:.1}%", change_ratio * 100.0);
|
||||
}
|
||||
Reference in New Issue
Block a user