#[cfg(test)]
use crate::sharding::coordinator::ShardAssignment;
#[cfg(test)]
use crate::sharding::distribution::compute_assignments;
#[cfg(test)]
use std::collections::HashMap;
#[test]
fn test_reproduce_user_overload_issue() {
let mut old_assignments = HashMap::new();
old_assignments.insert(
0,
ShardAssignment {
shard_id: 0,
primary_node: "6745".to_string(),
replica_nodes: vec!["6746".to_string()],
},
);
old_assignments.insert(
1,
ShardAssignment {
shard_id: 1,
primary_node: "6746".to_string(),
replica_nodes: vec!["6747".to_string()],
},
);
old_assignments.insert(
2,
ShardAssignment {
shard_id: 2,
primary_node: "6747".to_string(),
replica_nodes: vec!["6745".to_string()],
},
);
let nodes = vec!["6746".to_string(), "6747".to_string(), "6748".to_string()];
let assignments = compute_assignments(&nodes, 3, 2, Some(&old_assignments)).unwrap();
for id in 0..3 {
let a = &assignments[&id];
println!(
"Shard {}: P={}, R={:?}",
id, a.primary_node, a.replica_nodes
);
}
let mut primary_counts = HashMap::new();
for a in assignments.values() {
*primary_counts.entry(a.primary_node.clone()).or_insert(0) += 1;
}
assert!(primary_counts.get("6748").copied().unwrap_or(0) > 0,
"Node 6748 should be recruited as primary for one of the shards to balance load. 6748 count: {:?}",
primary_counts.get("6748"));
}