this repo has no description
1use std::sync::Arc;
2use dashmap::DashMap;
3use crate::{NodeId, Result, GigabrainError};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
6pub struct ShardId(pub u32);
7
8pub struct ShardManager {
9 shard_count: u32,
10 node_to_shard: Arc<DashMap<NodeId, ShardId>>,
11}
12
13impl ShardManager {
14 pub fn new(shard_count: u32) -> Self {
15 Self {
16 shard_count,
17 node_to_shard: Arc::new(DashMap::new()),
18 }
19 }
20
21 pub fn get_shard_for_node(&self, node_id: NodeId) -> ShardId {
22 if let Some(shard) = self.node_to_shard.get(&node_id) {
23 return *shard;
24 }
25
26 let shard_id = ShardId((node_id.0 % self.shard_count as u64) as u32);
27 self.node_to_shard.insert(node_id, shard_id);
28 shard_id
29 }
30
31 pub fn rebalance(&self) -> Result<()> {
32 Err(GigabrainError::Distributed("Rebalancing not implemented".to_string()))
33 }
34}