sz_orm_sharding/
lib.rs

1//! # SZ-ORM Sharding — 分片路由
2//!
3//! 提供基于 FNV-1a + fmix64 终结化的确定性哈希与一致性哈希环分片路由，
4//! 保证跨进程/重启后同一 key 的路由结果一致，避免相似 key 聚集。
5//!
6//! ## 主要模块
7//!
8//! - [`enhanced`] — 增强分片能力（虚拟节点等）
9//! - [`routing`] — 分片键提取器（`ShardKeyExtractor` 等）
10//! - [`scatter`] — 跨分片聚合（Scatter-Gather）
11//! - [`cross_shard_tx`] — 跨分片事务协调（2PC / Best Effort）
12
13use serde::{Deserialize, Serialize};
14use std::collections::{HashMap, HashSet};
15use std::error::Error;
16use std::fmt;
17
18pub mod cross_shard_tx;
19pub mod enhanced;
20pub mod routing;
21pub mod scatter;
22
23// 顶层再导出常用类型，方便用户直接 `use sz_orm_sharding::*`
24pub use cross_shard_tx::{
25    ShardParticipant, ShardTransactionCoordinator, ShardTxError, ShardTxResult,
26};
27pub use routing::{CompositeKeyExtractor, FieldExtractor, ShardKeyExtractor};
28pub use scatter::ScatterGather;
29
30/// FNV-1a 64-bit 确定性哈希函数（带 MurmurHash3 fmix64 终结化）
31///
32/// 用于分片路由，保证跨进程/重启后同一 key 的哈希结果一致。
33/// 不依赖任何随机种子，避免 `DefaultHasher`（基于 `RandomState`）的不确定性。
34///
35/// 注意：纯 FNV-1a 对短字符串的雪崩特性较弱，相似前缀的 key（如 `key_0`、`key_1`）
36/// 哈希值高度相关，会导致一致性哈希环上分布严重不均。追加 fmix64 终结化步骤
37/// 打破这种结构相关性，使哈希值在 64-bit 空间中近似均匀分布。
38fn fnv1a_hash(data: &str) -> u64 {
39    const FNV_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
40    const FNV_PRIME: u64 = 0x100000001b3;
41    let mut hash = FNV_OFFSET_BASIS;
42    for &byte in data.as_bytes() {
43        hash ^= byte as u64;
44        hash = hash.wrapping_mul(FNV_PRIME);
45    }
46    // MurmurHash3 fmix64 终结化：保证良好雪崩特性，避免相似 key 聚集
47    hash ^= hash >> 33;
48    hash = hash.wrapping_mul(0xff51afd7ed558ccd);
49    hash ^= hash >> 33;
50    hash = hash.wrapping_mul(0xc4ceb9fe1a85ec53);
51    hash ^= hash >> 33;
52    hash
53}
54
55/// 分片策略
56///
57/// 注意：v0.3.0 起扩展为非 `Copy` 枚举（新增 `Enum`/`List`/`Directory`/`Composite`
58/// 携带数据的变体）。`Hash`/`Range`/`Date` 三个原始变体的路由行为保持向后兼容。
59#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
60pub enum ShardingStrategy {
61    /// 哈希分片：对 key 做哈希后取模选择 shard
62    Hash,
63    /// 范围分片：按 key 的字节值范围选择 shard
64    Range,
65    /// 日期分片：按 key 中包含的日期信息（YYYY-MM-DD）选择 shard
66    Date,
67    /// 枚举分片：显式 key → shard 映射，未匹配走默认 shard
68    Enum {
69        /// 显式映射表
70        mapping: HashMap<String, String>,
71        /// 未匹配时的默认 shard
72        default: Option<String>,
73    },
74    /// 列表分片：key 在预定义集合中则路由到 target，否则走默认
75    List {
76        /// 预定义 key 集合
77        keys: HashSet<String>,
78        /// 命中时路由到的目标 shard
79        target: String,
80        /// 未命中时的默认 shard
81        default: Option<String>,
82    },
83    /// 目录分片：动态查询路由表（key → shard）
84    Directory {
85        /// 动态路由表
86        table: HashMap<String, String>,
87    },
88    /// 复合分片：先按 primary 路由得到 group，再用 secondary 对 "group:key" 二级路由
89    Composite {
90        /// 一级策略（决定 group）
91        primary: Box<ShardingStrategy>,
92        /// 一级策略使用的 shard 列表（即 group 标签集合）
93        primary_shards: Vec<String>,
94        /// 二级策略（在 group 内路由）
95        secondary: Box<ShardingStrategy>,
96        /// 二级策略使用的 shard 列表（最终 shard）
97        secondary_shards: Vec<String>,
98    },
99}
100
101/// 分片路由错误
102#[derive(Debug, Clone, PartialEq, Eq)]
103pub enum ShardingError {
104    /// 未配置任何 shard，无法路由
105    NoShardsConfigured,
106    /// `Enum`/`List`/`Directory` 策略未匹配到 key，且无默认 shard
107    NoMappingForKey(String),
108    /// 工作线程 panic（用于 `ScatterGather` 并行场景）
109    ThreadPanic,
110}
111
112impl fmt::Display for ShardingError {
113    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
114        match self {
115            ShardingError::NoShardsConfigured => {
116                write!(f, "ShardingRouter has no shards configured")
117            }
118            ShardingError::NoMappingForKey(key) => write!(f, "no mapping for key: {}", key),
119            ShardingError::ThreadPanic => write!(f, "worker thread panicked"),
120        }
121    }
122}
123
124impl Error for ShardingError {}
125
126/// 分片路由器
127///
128/// 根据 `ShardingStrategy` 将 key 路由到对应的 shard。
129/// v0.3.0 起支持 `Enum`/`List`/`Directory`/`Composite` 等新策略。
130pub struct ShardingRouter {
131    strategy: ShardingStrategy,
132    /// 仅 Hash/Range/Date 使用；其他策略自带数据，忽略此字段
133    shards: Vec<String>,
134}
135
136impl ShardingRouter {
137    /// 创建路由器（兼容旧 API：传入策略与 shard 列表）
138    pub fn new(strategy: ShardingStrategy, shards: Vec<&str>) -> Self {
139        Self {
140            strategy,
141            shards: shards.into_iter().map(|s| s.to_string()).collect(),
142        }
143    }
144
145    /// 构造枚举分片路由器
146    pub fn new_enum(mapping: HashMap<String, String>, default: Option<String>) -> Self {
147        Self {
148            strategy: ShardingStrategy::Enum { mapping, default },
149            shards: vec![],
150        }
151    }
152
153    /// 构造列表分片路由器
154    pub fn new_list(keys: HashSet<String>, target: String, default: Option<String>) -> Self {
155        Self {
156            strategy: ShardingStrategy::List {
157                keys,
158                target,
159                default,
160            },
161            shards: vec![],
162        }
163    }
164
165    /// 构造目录分片路由器
166    pub fn new_directory(table: HashMap<String, String>) -> Self {
167        Self {
168            strategy: ShardingStrategy::Directory { table },
169            shards: vec![],
170        }
171    }
172
173    /// 构造复合分片路由器
174    pub fn new_composite(
175        primary: ShardingStrategy,
176        primary_shards: Vec<String>,
177        secondary: ShardingStrategy,
178        secondary_shards: Vec<String>,
179    ) -> Self {
180        Self {
181            strategy: ShardingStrategy::Composite {
182                primary: Box::new(primary),
183                primary_shards,
184                secondary: Box::new(secondary),
185                secondary_shards,
186            },
187            shards: vec![],
188        }
189    }
190
191    /// 根据 key 路由到对应的 shard
192    ///
193    /// # Errors
194    ///
195    /// - Hash/Range/Date 策略下 shards 为空时返回 [`ShardingError::NoShardsConfigured`]
196    /// - Enum/List/Directory 策略未匹配且无默认时返回 [`ShardingError::NoMappingForKey`]
197    pub fn route(&self, key: &str) -> Result<&str, ShardingError> {
198        route_strategy(&self.strategy, &self.shards, key)
199    }
200
201    /// 通过数据对象 + 提取器路由：先从 `data` 提取 key，再 `route(key)`
202    ///
203    /// # Errors
204    ///
205    /// 提取失败或路由失败时返回对应 [`ShardingError`]。
206    pub fn route_by_data(
207        &self,
208        data: &dyn std::any::Any,
209        extractor: &dyn ShardKeyExtractor,
210    ) -> Result<&str, ShardingError> {
211        let key = extractor.extract(data)?;
212        self.route(&key)
213    }
214
215    /// 返回所有 shard（用于广播查询；仅 Hash/Range/Date 有效）
216    pub fn query_all(&self) -> &[String] {
217        &self.shards
218    }
219
220    /// 返回当前策略（克隆）
221    pub fn strategy(&self) -> ShardingStrategy {
222        self.strategy.clone()
223    }
224
225    /// 返回 shard 数量
226    pub fn shard_count(&self) -> usize {
227        self.shards.len()
228    }
229}
230
231/// 通用路由分发：根据策略选择 shard
232///
233/// 作为自由函数实现，便于 `Composite` 递归调用时复用同一套逻辑。
234/// 输出生命周期 `'a` 绑定到 `strategy` 与 `shards`（结果借用其中之一），
235/// 与 `key` 的生命周期无关。
236fn route_strategy<'a>(
237    strategy: &'a ShardingStrategy,
238    shards: &'a [String],
239    key: &str,
240) -> Result<&'a str, ShardingError> {
241    match strategy {
242        ShardingStrategy::Hash => {
243            if shards.is_empty() {
244                return Err(ShardingError::NoShardsConfigured);
245            }
246            Ok(route_hash(shards, key))
247        }
248        ShardingStrategy::Range => {
249            if shards.is_empty() {
250                return Err(ShardingError::NoShardsConfigured);
251            }
252            Ok(route_range(shards, key))
253        }
254        ShardingStrategy::Date => {
255            if shards.is_empty() {
256                return Err(ShardingError::NoShardsConfigured);
257            }
258            Ok(route_date(shards, key))
259        }
260        ShardingStrategy::Enum { mapping, default } => {
261            if let Some(shard) = mapping.get(key) {
262                Ok(shard.as_str())
263            } else if let Some(d) = default {
264                Ok(d.as_str())
265            } else {
266                Err(ShardingError::NoMappingForKey(key.to_string()))
267            }
268        }
269        ShardingStrategy::List {
270            keys,
271            target,
272            default,
273        } => {
274            if keys.contains(key) {
275                Ok(target.as_str())
276            } else if let Some(d) = default {
277                Ok(d.as_str())
278            } else {
279                Err(ShardingError::NoMappingForKey(key.to_string()))
280            }
281        }
282        ShardingStrategy::Directory { table } => table
283            .get(key)
284            .map(|s| s.as_str())
285            .ok_or_else(|| ShardingError::NoMappingForKey(key.to_string())),
286        ShardingStrategy::Composite {
287            primary,
288            primary_shards,
289            secondary,
290            secondary_shards,
291        } => {
292            // 一级路由得到 group 标签
293            let group = route_strategy(primary, primary_shards, key)?;
294            // 用 "group:key" 作为二级 key，让二级策略在 group 命名空间内路由
295            let composite_key = format!("{}:{}", group, key);
296            route_strategy(secondary, secondary_shards, &composite_key)
297        }
298    }
299}
300
301/// 哈希路由：对 key 做哈希后取模选择 shard
302///
303/// 返回值从 `shards` 借用（生命周期 `'a`），与 `key` 无关。
304fn route_hash<'a>(shards: &'a [String], key: &str) -> &'a str {
305    let hash = fnv1a_hash(key);
306    let idx = (hash as usize) % shards.len();
307    &shards[idx]
308}
309
310/// 范围路由：按 key 的首字节将 keyspace [0, 256) 均分到各 shard
311///
312/// 返回值从 `shards` 借用（生命周期 `'a`），与 `key` 无关。
313fn route_range<'a>(shards: &'a [String], key: &str) -> &'a str {
314    let first_byte = key.bytes().next().unwrap_or(0) as usize;
315    let idx = (first_byte * shards.len()) / 256;
316    &shards[idx.min(shards.len() - 1)]
317}
318
319/// 日期路由：按 key 中包含的日期信息（YYYY-MM-DD）的"日"取模选择 shard
320///
321/// 返回值从 `shards` 借用（生命周期 `'a`），与 `key` 无关。
322fn route_date<'a>(shards: &'a [String], key: &str) -> &'a str {
323    if let Some(date) = extract_date(key) {
324        // 用日期中的"日"（day of month）取模
325        if let Some(day) = date.get(8..10).and_then(|s| s.parse::<usize>().ok()) {
326            if day >= 1 {
327                let idx = (day - 1) % shards.len();
328                return &shards[idx];
329            }
330        }
331        // 日期解析失败，回退到日期字符串的哈希
332        let hash = fnv1a_hash(&date);
333        let idx = (hash as usize) % shards.len();
334        return &shards[idx];
335    }
336    // 没有日期信息，回退到 key 整体哈希
337    let hash = fnv1a_hash(key);
338    let idx = (hash as usize) % shards.len();
339    &shards[idx]
340}
341
342/// 从字符串中提取 YYYY-MM-DD 格式的日期
343fn extract_date(key: &str) -> Option<String> {
344    let bytes = key.as_bytes();
345    if bytes.len() < 10 {
346        return None;
347    }
348    for i in 0..=bytes.len() - 10 {
349        if is_digit(bytes[i])
350            && is_digit(bytes[i + 1])
351            && is_digit(bytes[i + 2])
352            && is_digit(bytes[i + 3])
353            && bytes[i + 4] == b'-'
354            && is_digit(bytes[i + 5])
355            && is_digit(bytes[i + 6])
356            && bytes[i + 7] == b'-'
357            && is_digit(bytes[i + 8])
358            && is_digit(bytes[i + 9])
359        {
360            return String::from_utf8(bytes[i..i + 10].to_vec()).ok();
361        }
362    }
363    None
364}
365
366fn is_digit(b: u8) -> bool {
367    b.is_ascii_digit()
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373    use std::collections::HashSet;
374
375    // --- 基础测试 ---
376
377    #[test]
378    fn test_router_creation() {
379        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["shard0", "shard1"]);
380        assert_eq!(router.shard_count(), 2);
381        assert_eq!(router.strategy(), ShardingStrategy::Hash);
382    }
383
384    #[test]
385    fn test_query_all() {
386        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s1", "s2", "s3"]);
387        assert_eq!(router.query_all().len(), 3);
388        assert_eq!(router.query_all()[0], "s1");
389        assert_eq!(router.query_all()[2], "s3");
390    }
391
392    #[test]
393    fn test_empty_shards_returns_error() {
394        let router = ShardingRouter::new(ShardingStrategy::Hash, vec![]);
395        let result = router.route("any_key");
396        assert!(matches!(result, Err(ShardingError::NoShardsConfigured)));
397        if let Err(err) = result {
398            let msg = format!("{}", err);
399            assert!(
400                msg.contains("no shards configured"),
401                "error message should mention empty shards, got: {}",
402                msg
403            );
404        }
405    }
406
407    #[test]
408    fn test_single_shard_always_returns_it() {
409        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["only"]);
410        assert_eq!(router.route("any_key").unwrap(), "only");
411        assert_eq!(router.route("different").unwrap(), "only");
412
413        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["only"]);
414        assert_eq!(router.route("any_key").unwrap(), "only");
415
416        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["only"]);
417        assert_eq!(router.route("2026-07-18").unwrap(), "only");
418    }
419
420    // --- Hash 策略测试 ---
421
422    #[test]
423    fn test_hash_deterministic() {
424        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1", "s2"]);
425        // 同一 key 应总是路由到同一 shard
426        let first = router.route("user:123").unwrap();
427        for _ in 0..5 {
428            assert_eq!(
429                router.route("user:123").unwrap(),
430                first,
431                "Hash 路由应确定性"
432            );
433        }
434    }
435
436    #[test]
437    fn test_hash_different_keys_distribute() {
438        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1", "s2", "s3"]);
439        // 大量不同 key 应分布到多个 shard
440        let mut shards_hit = HashSet::new();
441        for i in 0..100 {
442            let key = format!("key_{}", i);
443            shards_hit.insert(router.route(&key).unwrap().to_string());
444        }
445        assert!(
446            shards_hit.len() >= 2,
447            "Hash 策略在 100 个不同 key 上应至少命中 2 个 shard，实际: {}",
448            shards_hit.len()
449        );
450    }
451
452    #[test]
453    fn test_hash_same_key_same_shard() {
454        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1", "s2"]);
455        let r1 = router.route("consistent_key").unwrap();
456        let r2 = router.route("consistent_key").unwrap();
457        let r3 = router.route("consistent_key").unwrap();
458        assert_eq!(r1, r2);
459        assert_eq!(r2, r3);
460    }
461
462    #[test]
463    fn test_hash_empty_key() {
464        let router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1"]);
465        let shard = router.route("").unwrap();
466        // 空 key 也应路由到某个有效 shard
467        assert!(shard == "s0" || shard == "s1");
468    }
469
470    // --- Range 策略测试 ---
471
472    #[test]
473    fn test_range_ascii_vs_non_ascii() {
474        // 2 个 shard：首字节 0-127 -> s0, 128+ -> s1
475        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s0", "s1"]);
476        // ASCII 字符首字节 0-127，路由到 s0
477        assert_eq!(router.route("Hello").unwrap(), "s0");
478        assert_eq!(router.route("world").unwrap(), "s0");
479        assert_eq!(router.route("A").unwrap(), "s0");
480        assert_eq!(router.route("a").unwrap(), "s0");
481        // 非 ASCII 字符首字节 >= 194，路由到 s1
482        assert_eq!(router.route("你好").unwrap(), "s1");
483        assert_eq!(router.route("é").unwrap(), "s1");
484    }
485
486    #[test]
487    fn test_range_different_keys_hit_different_shards() {
488        // 3 个 shard，构造能命中所有 shard 的 key
489        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s0", "s1", "s2"]);
490        let mut shards_hit = HashSet::new();
491        // 'A' = 65 -> (65*3)/256 = 0 -> s0
492        shards_hit.insert(router.route("A").unwrap().to_string());
493        // 'a' = 97 -> (97*3)/256 = 1 -> s1
494        shards_hit.insert(router.route("a").unwrap().to_string());
495        // 'é' 首字节 195 -> (195*3)/256 = 2 -> s2
496        shards_hit.insert(router.route("é").unwrap().to_string());
497        assert_eq!(
498            shards_hit.len(),
499            3,
500            "Range 策略应能命中所有 3 个 shard，实际: {:?}",
501            shards_hit
502        );
503    }
504
505    #[test]
506    fn test_range_deterministic() {
507        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s0", "s1", "s2"]);
508        let first = router.route("hello").unwrap();
509        assert_eq!(router.route("hello").unwrap(), first);
510        assert_eq!(router.route("hello").unwrap(), first);
511    }
512
513    #[test]
514    fn test_range_empty_key_uses_zero_byte() {
515        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s0", "s1"]);
516        // 空 key 的首字节视为 0，应该路由到 s0
517        assert_eq!(router.route("").unwrap(), "s0");
518    }
519
520    #[test]
521    fn test_range_keys_with_similar_prefixes_cluster() {
522        // 相似前缀的 key 应路由到相同 shard（Range 的核心特性）
523        let router = ShardingRouter::new(ShardingStrategy::Range, vec!["s0", "s1"]);
524        let shard1 = router.route("user:123").unwrap();
525        let shard2 = router.route("user:456").unwrap();
526        let shard3 = router.route("user:789").unwrap();
527        assert_eq!(shard1, shard2);
528        assert_eq!(shard2, shard3);
529    }
530
531    // --- Date 策略测试 ---
532
533    #[test]
534    fn test_date_day_based_routing() {
535        // 3 个 shard，day 1..31 取模 3
536        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
537        // day=1 -> (1-1)%3 = 0 -> s0
538        assert_eq!(router.route("2026-07-01").unwrap(), "s0");
539        // day=2 -> (2-1)%3 = 1 -> s1
540        assert_eq!(router.route("2026-07-02").unwrap(), "s1");
541        // day=3 -> (3-1)%3 = 2 -> s2
542        assert_eq!(router.route("2026-07-03").unwrap(), "s2");
543        // day=4 -> (4-1)%3 = 0 -> s0
544        assert_eq!(router.route("2026-07-04").unwrap(), "s0");
545    }
546
547    #[test]
548    fn test_date_different_days_distribute() {
549        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2", "s3"]);
550        let mut shards_hit = HashSet::new();
551        for day in 1..=28 {
552            let key = format!("2026-07-{:02}", day);
553            shards_hit.insert(router.route(&key).unwrap().to_string());
554        }
555        // 28 天应分布到所有 4 个 shard
556        assert_eq!(
557            shards_hit.len(),
558            4,
559            "Date 策略 28 天应命中所有 4 个 shard，实际: {}",
560            shards_hit.len()
561        );
562    }
563
564    #[test]
565    fn test_date_extract_from_longer_key() {
566        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
567        // 日期嵌入在更长的 key 中
568        let shard1 = router.route("log:2026-07-15:entry1").unwrap();
569        let shard2 = router.route("2026-07-15").unwrap();
570        assert_eq!(shard1, shard2, "包含相同日期的 key 应路由到相同 shard");
571    }
572
573    #[test]
574    fn test_date_deterministic() {
575        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
576        let first = router.route("2026-07-18").unwrap();
577        assert_eq!(router.route("2026-07-18").unwrap(), first);
578    }
579
580    #[test]
581    fn test_date_no_date_falls_back_to_hash() {
582        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
583        // 没有日期信息的 key 应回退到哈希路由（仍返回有效 shard）
584        let shard = router.route("plain_key_without_date").unwrap();
585        assert!(shard == "s0" || shard == "s1" || shard == "s2");
586        // 且确定性
587        assert_eq!(router.route("plain_key_without_date").unwrap(), shard);
588    }
589
590    #[test]
591    fn test_date_different_months_same_day_same_shard() {
592        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
593        // 同一天不同月应路由到相同 shard（因为只看 day）
594        let july_15 = router.route("2026-07-15").unwrap();
595        let aug_15 = router.route("2026-08-15").unwrap();
596        assert_eq!(july_15, aug_15);
597    }
598
599    #[test]
600    fn test_date_invalid_date_falls_back() {
601        let router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1"]);
602        // "2026-00-00" 日为 00，无法解析为有效 day（parse::<usize> 得到 0，不满足 >= 1）
603        // 应回退到日期字符串的哈希
604        let shard = router.route("2026-00-00").unwrap();
605        assert!(shard == "s0" || shard == "s1");
606    }
607
608    // --- 跨策略测试 ---
609
610    #[test]
611    fn test_different_strategies_may_route_differently() {
612        let key = "2026-07-15";
613        let hash_router = ShardingRouter::new(ShardingStrategy::Hash, vec!["s0", "s1", "s2"]);
614        let date_router = ShardingRouter::new(ShardingStrategy::Date, vec!["s0", "s1", "s2"]);
615
616        // 不要求一定不同，但都应返回有效 shard
617        let hash_shard = hash_router.route(key).unwrap();
618        let date_shard = date_router.route(key).unwrap();
619        assert!(!hash_shard.is_empty());
620        assert!(!date_shard.is_empty());
621    }
622
623    // --- extract_date 单元测试 ---
624
625    #[test]
626    fn test_extract_date_pure_date() {
627        assert_eq!(extract_date("2026-07-18"), Some("2026-07-18".to_string()));
628        assert_eq!(extract_date("2025-01-01"), Some("2025-01-01".to_string()));
629    }
630
631    #[test]
632    fn test_extract_date_embedded() {
633        assert_eq!(
634            extract_date("log:2026-07-18:entry"),
635            Some("2026-07-18".to_string())
636        );
637    }
638
639    #[test]
640    fn test_extract_date_no_date() {
641        assert_eq!(extract_date("no date here"), None);
642        assert_eq!(extract_date("2026/07/18"), None);
643        assert_eq!(extract_date(""), None);
644        assert_eq!(extract_date("short"), None);
645    }
646
647    #[test]
648    fn test_extract_date_invalid_format() {
649        assert_eq!(extract_date("2026-7-18"), None); // 月需要 2 位
650        assert_eq!(extract_date("2026-07-8"), None); // 日需要 2 位
651        assert_eq!(extract_date("abcd-07-18"), None); // 年需为数字
652    }
653
654    // ==================== v0.3.0 新增策略测试 ====================
655
656    // --- Enum 策略测试 ---
657
658    #[test]
659    fn test_enum_route_hit() {
660        let mut mapping = HashMap::new();
661        mapping.insert("cn".to_string(), "shard_cn".to_string());
662        mapping.insert("us".to_string(), "shard_us".to_string());
663        mapping.insert("eu".to_string(), "shard_eu".to_string());
664        let router = ShardingRouter::new_enum(mapping, None);
665        assert_eq!(router.route("cn").unwrap(), "shard_cn");
666        assert_eq!(router.route("us").unwrap(), "shard_us");
667        assert_eq!(router.route("eu").unwrap(), "shard_eu");
668    }
669
670    #[test]
671    fn test_enum_route_miss_with_default() {
672        let mut mapping = HashMap::new();
673        mapping.insert("cn".to_string(), "shard_cn".to_string());
674        let router = ShardingRouter::new_enum(mapping, Some("shard_default".to_string()));
675        assert_eq!(router.route("unknown").unwrap(), "shard_default");
676        // 命中映射的仍返回映射值
677        assert_eq!(router.route("cn").unwrap(), "shard_cn");
678    }
679
680    #[test]
681    fn test_enum_route_miss_no_default_errors() {
682        let router = ShardingRouter::new_enum(HashMap::new(), None);
683        let result = router.route("unknown");
684        assert!(matches!(result, Err(ShardingError::NoMappingForKey(_))));
685        if let Err(ShardingError::NoMappingForKey(key)) = result {
686            assert_eq!(key, "unknown");
687        } else {
688            panic!("expected NoMappingForKey");
689        }
690    }
691
692    #[test]
693    fn test_enum_route_deterministic() {
694        let mut mapping = HashMap::new();
695        mapping.insert("k1".to_string(), "s_a".to_string());
696        let router = ShardingRouter::new_enum(mapping, Some("s_def".to_string()));
697        let r1 = router.route("k1").unwrap();
698        let r2 = router.route("k1").unwrap();
699        assert_eq!(r1, r2);
700        assert_eq!(router.route("k2").unwrap(), "s_def");
701    }
702
703    // --- List 策略测试 ---
704
705    #[test]
706    fn test_list_route_hit() {
707        let mut keys = HashSet::new();
708        keys.insert("vip1".to_string());
709        keys.insert("vip2".to_string());
710        keys.insert("vip3".to_string());
711        let router = ShardingRouter::new_list(keys, "vip_shard".to_string(), None);
712        assert_eq!(router.route("vip1").unwrap(), "vip_shard");
713        assert_eq!(router.route("vip2").unwrap(), "vip_shard");
714        assert_eq!(router.route("vip3").unwrap(), "vip_shard");
715    }
716
717    #[test]
718    fn test_list_route_miss_with_default() {
719        let keys = HashSet::new();
720        let router = ShardingRouter::new_list(
721            keys,
722            "vip_shard".to_string(),
723            Some("normal_shard".to_string()),
724        );
725        // 命中列表的 key 路由到 target
726        assert_eq!(router.route("any_non_listed").unwrap(), "normal_shard");
727    }
728
729    #[test]
730    fn test_list_route_miss_no_default_errors() {
731        let router = ShardingRouter::new_list(HashSet::new(), "vip_shard".to_string(), None);
732        let result = router.route("unknown");
733        assert!(matches!(result, Err(ShardingError::NoMappingForKey(_))));
734    }
735
736    #[test]
737    fn test_list_route_with_members_and_default() {
738        let mut keys = HashSet::new();
739        keys.insert("gold".to_string());
740        let router = ShardingRouter::new_list(
741            keys,
742            "premium_shard".to_string(),
743            Some("standard_shard".to_string()),
744        );
745        assert_eq!(router.route("gold").unwrap(), "premium_shard");
746        assert_eq!(router.route("silver").unwrap(), "standard_shard");
747    }
748
749    // --- Directory 策略测试 ---
750
751    #[test]
752    fn test_directory_route_hit() {
753        let mut table = HashMap::new();
754        table.insert("user:1".to_string(), "dir_shard_a".to_string());
755        table.insert("user:2".to_string(), "dir_shard_b".to_string());
756        let router = ShardingRouter::new_directory(table);
757        assert_eq!(router.route("user:1").unwrap(), "dir_shard_a");
758        assert_eq!(router.route("user:2").unwrap(), "dir_shard_b");
759    }
760
761    #[test]
762    fn test_directory_route_miss_errors() {
763        let router = ShardingRouter::new_directory(HashMap::new());
764        assert!(matches!(
765            router.route("missing"),
766            Err(ShardingError::NoMappingForKey(_))
767        ));
768    }
769
770    #[test]
771    fn test_directory_route_deterministic() {
772        let mut table = HashMap::new();
773        table.insert("k".to_string(), "v".to_string());
774        let router = ShardingRouter::new_directory(table);
775        let r1 = router.route("k").unwrap();
776        let r2 = router.route("k").unwrap();
777        assert_eq!(r1, r2);
778        assert_eq!(r1, "v");
779    }
780
781    // --- Composite 策略测试 ---
782
783    #[test]
784    fn test_composite_route_basic() {
785        // 一级：Hash 在 [g0, g1] 上路由得 group
786        // 二级：Hash 在 [s0, s1, s2] 上路由得最终 shard
787        let router = ShardingRouter::new_composite(
788            ShardingStrategy::Hash,
789            vec!["g0".to_string(), "g1".to_string()],
790            ShardingStrategy::Hash,
791            vec!["s0".to_string(), "s1".to_string(), "s2".to_string()],
792        );
793        let result = router.route("user:123").unwrap();
794        assert!(
795            result == "s0" || result == "s1" || result == "s2",
796            "composite result should be in secondary shards, got {}",
797            result
798        );
799    }
800
801    #[test]
802    fn test_composite_route_deterministic() {
803        let router = ShardingRouter::new_composite(
804            ShardingStrategy::Hash,
805            vec!["g0".to_string(), "g1".to_string()],
806            ShardingStrategy::Hash,
807            vec!["s0".to_string(), "s1".to_string()],
808        );
809        let r1 = router.route("user:123").unwrap();
810        for _ in 0..5 {
811            assert_eq!(router.route("user:123").unwrap(), r1);
812        }
813    }
814
815    #[test]
816    fn test_composite_uses_group_in_secondary_key() {
817        // 构造两个不同 group 的场景：一级用 Enum 强制分组
818        // groupA 的所有 key 经二级 Hash 在 [s0, s1] 路由
819        // groupB 的所有 key 经二级 Hash 在 [s0, s1] 路由，但二级 key 含不同 group 前缀
820        let mut mapping = HashMap::new();
821        mapping.insert("a".to_string(), "grpA".to_string());
822        mapping.insert("b".to_string(), "grpB".to_string());
823        let router = ShardingRouter::new_composite(
824            ShardingStrategy::Enum {
825                mapping,
826                default: None,
827            },
828            vec!["grpA".to_string(), "grpB".to_string()], // primary_shards（仅占位，Enum 不使用）
829            ShardingStrategy::Hash,
830            vec!["s0".to_string(), "s1".to_string()],
831        );
832        // key "a" 与 "b" 走不同 group，二级用 "grpA:a" / "grpB:b" 路由
833        let ra = router.route("a").unwrap();
834        let rb = router.route("b").unwrap();
835        // 都应路由到有效 shard
836        assert!(ra == "s0" || ra == "s1");
837        assert!(rb == "s0" || rb == "s1");
838    }
839
840    #[test]
841    fn test_composite_primary_empty_shards_errors() {
842        // primary 用 Hash 但 primary_shards 为空 → NoShardsConfigured
843        let router = ShardingRouter::new_composite(
844            ShardingStrategy::Hash,
845            vec![],
846            ShardingStrategy::Hash,
847            vec!["s0".to_string()],
848        );
849        assert!(matches!(
850            router.route("k"),
851            Err(ShardingError::NoShardsConfigured)
852        ));
853    }
854
855    #[test]
856    fn test_composite_secondary_empty_shards_errors() {
857        // secondary 用 Hash 但 secondary_shards 为空 → NoShardsConfigured
858        let router = ShardingRouter::new_composite(
859            ShardingStrategy::Hash,
860            vec!["g0".to_string()],
861            ShardingStrategy::Hash,
862            vec![],
863        );
864        assert!(matches!(
865            router.route("k"),
866            Err(ShardingError::NoShardsConfigured)
867        ));
868    }
869
870    // --- ShardingError 新增变体测试 ---
871
872    #[test]
873    fn test_sharding_error_no_mapping_display() {
874        let err = ShardingError::NoMappingForKey("k1".to_string());
875        let msg = format!("{}", err);
876        assert!(msg.contains("no mapping for key: k1"));
877    }
878
879    #[test]
880    fn test_sharding_error_thread_panic_display() {
881        let err = ShardingError::ThreadPanic;
882        let msg = format!("{}", err);
883        assert!(msg.contains("panicked"));
884    }
885}
sz_orm_sharding/lib.rs

sz_orm_sharding/
lib.rs
