143 lines
4.8 KiB
Rust
143 lines
4.8 KiB
Rust
use flaredb_proto::pdpb::{Region, Store};
|
|
use std::collections::HashMap;
|
|
use std::sync::{Arc, Mutex};
|
|
|
|
#[derive(Clone)]
|
|
pub struct Cluster {
|
|
inner: Arc<Mutex<ClusterState>>,
|
|
}
|
|
|
|
struct ClusterState {
|
|
stores: HashMap<u64, Store>,
|
|
regions: Vec<Region>, // Sorted by start_key? Or just linear scan for MVP
|
|
region_leaders: HashMap<u64, u64>, // region_id -> store_id
|
|
next_store_id: u64,
|
|
next_region_id: u64,
|
|
}
|
|
|
|
impl Cluster {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
inner: Arc::new(Mutex::new(ClusterState {
|
|
stores: HashMap::new(),
|
|
regions: Vec::new(),
|
|
region_leaders: HashMap::new(),
|
|
next_store_id: 1,
|
|
next_region_id: 1,
|
|
})),
|
|
}
|
|
}
|
|
|
|
pub fn register_store(&self, addr: String, requested_id: Option<u64>) -> u64 {
|
|
let mut state = self.inner.lock().unwrap();
|
|
|
|
// Dedup check? For now, always new ID.
|
|
// In real system, check if addr exists.
|
|
for store in state.stores.values() {
|
|
if store.addr == addr {
|
|
return store.id;
|
|
}
|
|
}
|
|
|
|
let id = requested_id
|
|
.filter(|id| *id != 0 && !state.stores.contains_key(id))
|
|
.unwrap_or_else(|| {
|
|
while state.stores.contains_key(&state.next_store_id) {
|
|
state.next_store_id += 1;
|
|
}
|
|
state.next_store_id
|
|
});
|
|
state.next_store_id = state.next_store_id.max(id.saturating_add(1));
|
|
|
|
state.stores.insert(id, Store { id, addr });
|
|
|
|
// Bootstrap check: If we have 3 stores (Raft requires 3 for quorum), create Raft group
|
|
if state.stores.len() == 3 && state.regions.is_empty() {
|
|
let r1_id = state.next_region_id;
|
|
state.next_region_id += 1;
|
|
|
|
let r1 = Region {
|
|
id: r1_id,
|
|
start_key: vec![],
|
|
end_key: vec![], // All keys
|
|
peers: Vec::new(),
|
|
leader_id: 0,
|
|
};
|
|
|
|
state.regions.push(r1);
|
|
|
|
// Assign ALL 3 stores as peers (Leader = S1, but peers = [S1, S2, S3])
|
|
// We need to update Region definition to include Peers?
|
|
// Currently Region struct in pdpb doesn't have peers field (it's commented out).
|
|
// We can just assign leader for routing.
|
|
// Real Raft: Server needs to know peers.
|
|
// PD should return Peers in GetRegionResponse.
|
|
|
|
let ids: Vec<u64> = state.stores.keys().cloned().collect();
|
|
let mut sorted_ids = ids;
|
|
sorted_ids.sort();
|
|
|
|
state.region_leaders.insert(r1_id, sorted_ids[0]);
|
|
|
|
println!(
|
|
"Bootstrapped Raft Cluster: Region {} -> Leader S{}, Peers {:?}",
|
|
r1_id, sorted_ids[0], sorted_ids
|
|
);
|
|
} else if state.stores.len() == 1 && state.regions.is_empty() {
|
|
// Single node: one region ["", "")
|
|
let r1_id = state.next_region_id;
|
|
state.next_region_id += 1;
|
|
|
|
let r1 = Region {
|
|
id: r1_id,
|
|
start_key: vec![],
|
|
end_key: vec![],
|
|
peers: Vec::new(),
|
|
leader_id: 0,
|
|
};
|
|
state.regions.push(r1);
|
|
|
|
// Assign to this store
|
|
state.region_leaders.insert(r1_id, id);
|
|
println!("Bootstrapped Single Node: R1->S{}", id);
|
|
}
|
|
|
|
id
|
|
}
|
|
|
|
pub fn get_region_for_key(&self, key: &[u8]) -> Option<(Region, Store)> {
|
|
let state = self.inner.lock().unwrap();
|
|
|
|
for region in &state.regions {
|
|
let start_ok = region.start_key.is_empty() || key >= ®ion.start_key[..];
|
|
let end_ok = region.end_key.is_empty() || key < ®ion.end_key[..];
|
|
|
|
if start_ok && end_ok {
|
|
// Found region
|
|
if let Some(leader_id) = state.region_leaders.get(®ion.id) {
|
|
if let Some(store) = state.stores.get(leader_id) {
|
|
let mut r = region.clone();
|
|
r.leader_id = *leader_id;
|
|
r.peers = state.stores.keys().cloned().collect();
|
|
return Some((r, store.clone()));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
None
|
|
}
|
|
|
|
pub fn list_regions(&self) -> (Vec<Region>, Vec<Store>) {
|
|
let state = self.inner.lock().unwrap();
|
|
let mut regions = Vec::new();
|
|
for r in &state.regions {
|
|
let mut clone = r.clone();
|
|
clone.leader_id = *state.region_leaders.get(&r.id).unwrap_or(&0);
|
|
clone.peers = state.stores.keys().cloned().collect();
|
|
regions.push(clone);
|
|
}
|
|
let stores: Vec<Store> = state.stores.values().cloned().collect();
|
|
(regions, stores)
|
|
}
|
|
}
|