use crate::log_store::log_manager::LogConfig; use serde::{Deserialize, Serialize}; use ssz_derive::{Decode, Encode}; use std::{cell::RefCell, path::PathBuf, rc::Rc, str::FromStr}; pub const SHARD_CONFIG_KEY: &str = "shard_config"; #[derive(Clone)] pub struct Config { pub db_dir: PathBuf, pub log_config: LogConfig, } #[derive(Clone, Copy, Debug, Decode, Encode, Serialize, Deserialize, Eq, PartialEq)] #[serde(rename_all = "camelCase")] pub struct ShardConfig { pub shard_id: usize, pub num_shard: usize, } impl Default for ShardConfig { fn default() -> Self { Self { shard_id: 0, num_shard: 1, } } } impl FromStr for ShardConfig { type Err = String; fn from_str(s: &str) -> Result { let parts: Vec<&str> = s.trim().split('/').map(|s| s.trim()).collect(); if parts.len() != 2 { return Err("Incorrect format, expected like: '0 / 8'".into()); } let numerator = parts[0] .parse::() .map_err(|e| format!("Cannot parse shard position {:?}", e))?; let denominator = parts[1] .parse::() .map_err(|e| format!("Cannot parse shard position {:?}", e))?; Self::new(numerator, denominator) } } impl TryFrom> for ShardConfig { type Error = String; fn try_from(value: Option) -> Result { if let Some(position) = value { Self::from_str(&position) } else { Ok(Self::default()) } } } impl ShardConfig { pub fn new(id: usize, num: usize) -> Result { let config = ShardConfig { shard_id: id, num_shard: num, }; config.validate()?; Ok(config) } pub fn miner_shard_mask(&self) -> u64 { !(self.num_shard - 1) as u64 } pub fn miner_shard_id(&self) -> u64 { self.shard_id as u64 } pub fn validate(&self) -> Result<(), String> { if self.shard_id >= self.num_shard { return Err(format!( "Incorrect shard_id: expected [0, {}), actual {}", self.num_shard, self.shard_id )); } if self.num_shard == 0 { return Err("Shard num is 0".into()); } if !self.num_shard.is_power_of_two() { return Err(format!( "Incorrect shard group bytes: {}, should be power of two", self.num_shard )); } Ok(()) } pub fn in_range(&self, segment_index: u64) -> bool { segment_index as usize % self.num_shard == self.shard_id } pub fn next_segment_index(&self, current: usize, start_index: usize) -> usize { // `shift` should be 0 if `current` was returned by the same config. let shift = (start_index + current + self.num_shard - self.shard_id) % self.num_shard; current + self.num_shard - shift } /// Whether `self` intersect with the `other` shard config. pub fn intersect(&self, other: &ShardConfig) -> bool { let ShardConfig { num_shard: mut left_num_shard, shard_id: mut left_shard_id, } = self; let ShardConfig { num_shard: mut right_num_shard, shard_id: mut right_shard_id, } = other; while left_num_shard != right_num_shard { if left_num_shard < right_num_shard { right_num_shard /= 2; right_shard_id /= 2; } else { left_num_shard /= 2; left_shard_id /= 2; } } left_shard_id == right_shard_id } } struct ShardSegmentTreeNode { pub num_shard: usize, pub covered: bool, pub childs: [Option>>; 2], } impl ShardSegmentTreeNode { pub fn new(num_shard: usize) -> Self { ShardSegmentTreeNode { num_shard, covered: false, childs: [None, None], } } fn push_down(&mut self) { if self.childs[0].is_none() { for i in 0..2 { self.childs[i] = Some(Rc::new(RefCell::new(ShardSegmentTreeNode::new( self.num_shard << 1, )))); } } } fn update(&mut self) { let mut covered = true; for i in 0..2 { if let Some(child) = &self.childs[i] { covered = covered && child.borrow().covered; } } self.covered = covered; } pub fn insert(&mut self, num_shard: usize, shard_id: usize) { if self.covered { return; } if num_shard == self.num_shard { self.covered = true; return; } self.push_down(); if let Some(child) = &self.childs[shard_id % 2] { child.borrow_mut().insert(num_shard, shard_id >> 1); } self.update(); } } pub fn all_shards_available(shard_configs: Vec) -> bool { let mut root = ShardSegmentTreeNode::new(1); for shard_config in shard_configs.iter() { if shard_config.validate().is_err() { continue; } root.insert(shard_config.num_shard, shard_config.shard_id); if root.covered { return true; } } false } #[cfg(test)] mod tests { use crate::config::all_shards_available; use super::ShardConfig; fn new_config(id: usize, num: usize) -> ShardConfig { ShardConfig::new(id, num).unwrap() } #[test] fn test_all_shards_available() { assert!(all_shards_available(vec![ ShardConfig { shard_id: 3, num_shard: 8 }, ShardConfig { shard_id: 7, num_shard: 8 }, ShardConfig { shard_id: 0, num_shard: 4 }, ShardConfig { shard_id: 1, num_shard: 4 }, ShardConfig { shard_id: 0, num_shard: 2 }, ShardConfig { shard_id: 0, num_shard: 1 << 25 }, ])); assert!(!all_shards_available(vec![ ShardConfig { shard_id: 0, num_shard: 4 }, ShardConfig { shard_id: 1, num_shard: 4 }, ShardConfig { shard_id: 3, num_shard: 8 }, ShardConfig { shard_id: 0, num_shard: 2 }, ])); } #[test] fn test_shard_intersect() { // 1 shard assert!(new_config(0, 1).intersect(&new_config(0, 1))); // either is 1 shard assert!(new_config(0, 1).intersect(&new_config(0, 2))); assert!(new_config(0, 1).intersect(&new_config(1, 2))); assert!(new_config(0, 2).intersect(&new_config(0, 1))); assert!(new_config(1, 2).intersect(&new_config(0, 1))); // same shards assert!(!new_config(1, 4).intersect(&new_config(0, 4))); assert!(new_config(1, 4).intersect(&new_config(1, 4))); assert!(!new_config(1, 4).intersect(&new_config(2, 4))); assert!(!new_config(1, 4).intersect(&new_config(3, 4))); // left shards is less assert!(!new_config(1, 2).intersect(&new_config(0, 4))); assert!(!new_config(1, 2).intersect(&new_config(1, 4))); assert!(new_config(1, 2).intersect(&new_config(2, 4))); assert!(new_config(1, 2).intersect(&new_config(3, 4))); // right shards is less assert!(new_config(1, 4).intersect(&new_config(0, 2))); assert!(!new_config(1, 4).intersect(&new_config(1, 2))); assert!(!new_config(2, 4).intersect(&new_config(0, 2))); assert!(new_config(2, 4).intersect(&new_config(1, 2))); } }