Line | Count | Source |
1 | //! Anti-entropy reconciler module. | |
2 | //! | |
3 | //! Stub for plan §13.8 anti-entropy shard reconciler. | |
4 | //! Full implementation will follow the fingerprint → diff → repair pipeline. | |
5 | ||
6 | use serde::{Deserialize, Serialize}; | |
7 | ||
8 | use crate::migration::{MigrationConfig, MigrationError}; | |
9 | ||
10 | /// Anti-entropy configuration (plan §13.8). | |
11 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
12 | pub struct AntiEntropyConfig { | |
13 | pub enabled: bool, | |
14 | pub schedule_cron: String, | |
15 | pub shards_per_pass: u32, | |
16 | pub max_read_concurrency: u32, | |
17 | pub fingerprint_batch_size: u32, | |
18 | pub auto_repair: bool, | |
19 | pub updated_at_field: String, | |
20 | } | |
21 | ||
22 | impl Default for AntiEntropyConfig { | |
23 | 6 | fn default() -> Self { |
24 | 6 | Self { |
25 | 6 | enabled: true, |
26 | 6 | schedule_cron: "0 */6 * * *".to_string(), |
27 | 6 | shards_per_pass: 0, |
28 | 6 | max_read_concurrency: 2, |
29 | 6 | fingerprint_batch_size: 1000, |
30 | 6 | auto_repair: true, |
31 | 6 | updated_at_field: "_miroir_updated_at".to_string(), |
32 | 6 | } |
33 | 6 | } |
34 | } | |
35 | ||
36 | /// Validates that migration is safe given the anti-entropy configuration. | |
37 | /// Returns Ok(()) if safe, Err with a descriptive message if not. | |
38 | /// | |
39 | /// Hard refusal policy (plan §15 OP#1): skipping the delta pass while | |
40 | /// anti-entropy is disabled provides zero recovery path for documents | |
41 | /// written at the cutover boundary. Measured loss rate: ~2% of writes. | |
42 | /// This is a hard-coded policy, not a warning. | |
43 | 6 | pub fn validate_migration_safety( |
44 | 6 | ae_config: &AntiEntropyConfig, |
45 | 6 | migration_config: &MigrationConfig, |
46 | 6 | ) -> Result<(), MigrationError> { |
47 | 6 | if migration_config.skip_delta_pass && |
48 | 2 | return Err(MigrationError::UnsafeCutoverNoAntiEntropy); |
49 | 4 | } |
50 | 4 | Ok(()) |
51 | 6 | } |
52 | ||
53 | /// Generates a warning if anti-entropy is disabled during active migration. | |
54 | /// The caller should log this at warn level. | |
55 | /// | |
56 | /// Even with the delta pass enabled (which provides 0-loss cutover on its own), | |
57 | /// disabling anti-entropy means the delta pass is the sole safety mechanism. | |
58 | /// Operators should be aware of this reduced redundancy. | |
59 | 4 | pub fn migration_warning_if_ae_disabled(ae_enabled: bool) -> Option<String> { |
60 | 4 | if ae_enabled { |
61 | 2 | return None; |
62 | 2 | } |
63 | 2 | Some( |
64 | 2 | "Anti-entropy is disabled. Shard migration cutover relies on the delta pass \ |
65 | 2 | as the sole safety mechanism. Any bugs in the delta pass could lead to \ |
66 | 2 | data loss at the cutover boundary. Re-enable anti-entropy for defense-in-depth." |
67 | 2 | .to_string(), |
68 | 2 | ) |
69 | 4 | } |
70 | ||
71 | #[cfg(test)] | |
72 | mod tests { | |
73 | use super::*; | |
74 | ||
75 | #[test] | |
76 | 2 | fn test_validate_safe_with_delta_pass() { |
77 | 2 | let ae = AntiEntropyConfig { |
78 | 2 | enabled: false, |
79 | 2 | ..Default::default() |
80 | 2 | }; |
81 | 2 | let mc = MigrationConfig { |
82 | 2 | skip_delta_pass: false, |
83 | 2 | ..Default::default() |
84 | 2 | }; |
85 | 2 | assert!(validate_migration_safety(&ae, &mc).is_ok()); |
86 | 2 | } |
87 | ||
88 | #[test] | |
89 | 2 | fn test_validate_unsafe_without_anti_entropy() { |
90 | 2 | let ae = AntiEntropyConfig { |
91 | 2 | enabled: false, |
92 | 2 | ..Default::default() |
93 | 2 | }; |
94 | 2 | let mc = MigrationConfig { |
95 | 2 | skip_delta_pass: true, |
96 | 2 | anti_entropy_enabled: false, |
97 | 2 | ..Default::default() |
98 | 2 | }; |
99 | 2 | assert!(validate_migration_safety(&ae, &mc).is_err()); |
100 | 2 | } |
101 | ||
102 | #[test] | |
103 | 2 | fn test_validate_safe_with_anti_entropy_safety_net() { |
104 | 2 | let ae = AntiEntropyConfig { |
105 | 2 | enabled: true, |
106 | 2 | ..Default::default() |
107 | 2 | }; |
108 | 2 | let mc = MigrationConfig { |
109 | 2 | skip_delta_pass: true, |
110 | 2 | anti_entropy_enabled: true, |
111 | 2 | ..Default::default() |
112 | 2 | }; |
113 | 2 | assert!(validate_migration_safety(&ae, &mc).is_ok()); |
114 | 2 | } |
115 | ||
116 | #[test] | |
117 | 2 | fn test_warning_when_ae_disabled() { |
118 | 2 | assert!(migration_warning_if_ae_disabled(false).is_some()); |
119 | 2 | assert!(migration_warning_if_ae_disabled(true).is_none()); |
120 | 2 | } |
121 | } |
Line | Count | Source |
1 | //! Miroir configuration — plan §4 YAML schema with §13 advanced capabilities. | |
2 | ||
3 | mod advanced; | |
4 | mod error; | |
5 | mod load; | |
6 | mod validate; | |
7 | ||
8 | pub use error::ConfigError; | |
9 | ||
10 | use serde::{Deserialize, Serialize}; | |
11 | ||
12 | /// Top-level configuration matching plan §4 YAML schema under `miroir:`. | |
13 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
14 | #[serde(default)] | |
15 | pub struct MiroirConfig { | |
16 | // --- Secrets (env-var overrides) --- | |
17 | /// Client-facing API key. Env override: `MIROIR_MASTER_KEY`. | |
18 | pub master_key: String, | |
19 | /// Key Miroir uses on Meilisearch nodes. Env override: `MIROIR_NODE_MASTER_KEY`. | |
20 | pub node_master_key: String, | |
21 | ||
22 | // --- Core topology --- | |
23 | /// Total number of logical shards. | |
24 | pub shards: u32, | |
25 | /// Replication factor (intra-group replicas per shard). Production: 2. | |
26 | pub replication_factor: u32, | |
27 | /// Number of independent query pools. Default 1; production: 2. | |
28 | pub replica_groups: u32, | |
29 | ||
30 | // --- Sub-structs --- | |
31 | pub nodes: Vec<NodeConfig>, | |
32 | pub task_store: TaskStoreConfig, | |
33 | pub admin: AdminConfig, | |
34 | pub health: HealthConfig, | |
35 | pub scatter: ScatterConfig, | |
36 | pub rebalancer: RebalancerConfig, | |
37 | pub server: ServerConfig, | |
38 | pub connection_pool_per_node: ConnectionPoolConfig, | |
39 | pub task_registry: TaskRegistryConfig, | |
40 | ||
41 | // --- §13 advanced capabilities --- | |
42 | pub resharding: advanced::ReshardingConfig, | |
43 | pub hedging: advanced::HedgingConfig, | |
44 | pub replica_selection: advanced::ReplicaSelectionConfig, | |
45 | pub query_planner: advanced::QueryPlannerConfig, | |
46 | pub settings_broadcast: advanced::SettingsBroadcastConfig, | |
47 | pub settings_drift_check: advanced::SettingsDriftCheckConfig, | |
48 | pub session_pinning: advanced::SessionPinningConfig, | |
49 | pub aliases: advanced::AliasesConfig, | |
50 | pub anti_entropy: advanced::AntiEntropyConfig, | |
51 | pub dump_import: advanced::DumpImportConfig, | |
52 | pub idempotency: advanced::IdempotencyConfig, | |
53 | pub query_coalescing: advanced::QueryCoalescingConfig, | |
54 | pub multi_search: advanced::MultiSearchConfig, | |
55 | pub vector_search: advanced::VectorSearchConfig, | |
56 | pub cdc: advanced::CdcConfig, | |
57 | pub ttl: advanced::TtlConfig, | |
58 | pub tenant_affinity: advanced::TenantAffinityConfig, | |
59 | pub shadow: advanced::ShadowConfig, | |
60 | pub ilm: advanced::IlmConfig, | |
61 | pub canary_runner: advanced::CanaryRunnerConfig, | |
62 | pub explain: advanced::ExplainConfig, | |
63 | pub admin_ui: advanced::AdminUiConfig, | |
64 | pub search_ui: advanced::SearchUiConfig, | |
65 | ||
66 | // --- §14 horizontal scaling --- | |
67 | pub peer_discovery: PeerDiscoveryConfig, | |
68 | pub leader_election: LeaderElectionConfig, | |
69 | pub hpa: HpaConfig, | |
70 | } | |
71 | ||
72 | /// Convenience alias. | |
73 | pub type Config = MiroirConfig; | |
74 | ||
75 | impl Default for MiroirConfig { | |
76 | 32 | fn default() -> Self { |
77 | 32 | Self { |
78 | 32 | master_key: String::new(), |
79 | 32 | node_master_key: String::new(), |
80 | 32 | shards: 64, |
81 | 32 | replication_factor: 2, |
82 | 32 | replica_groups: 1, |
83 | 32 | nodes: Vec::new(), |
84 | 32 | task_store: TaskStoreConfig::default(), |
85 | 32 | admin: AdminConfig::default(), |
86 | 32 | health: HealthConfig::default(), |
87 | 32 | scatter: ScatterConfig::default(), |
88 | 32 | rebalancer: RebalancerConfig::default(), |
89 | 32 | server: ServerConfig::default(), |
90 | 32 | connection_pool_per_node: ConnectionPoolConfig::default(), |
91 | 32 | task_registry: TaskRegistryConfig::default(), |
92 | 32 | resharding: advanced::ReshardingConfig::default(), |
93 | 32 | hedging: advanced::HedgingConfig::default(), |
94 | 32 | replica_selection: advanced::ReplicaSelectionConfig::default(), |
95 | 32 | query_planner: advanced::QueryPlannerConfig::default(), |
96 | 32 | settings_broadcast: advanced::SettingsBroadcastConfig::default(), |
97 | 32 | settings_drift_check: advanced::SettingsDriftCheckConfig::default(), |
98 | 32 | session_pinning: advanced::SessionPinningConfig::default(), |
99 | 32 | aliases: advanced::AliasesConfig::default(), |
100 | 32 | anti_entropy: advanced::AntiEntropyConfig::default(), |
101 | 32 | dump_import: advanced::DumpImportConfig::default(), |
102 | 32 | idempotency: advanced::IdempotencyConfig::default(), |
103 | 32 | query_coalescing: advanced::QueryCoalescingConfig::default(), |
104 | 32 | multi_search: advanced::MultiSearchConfig::default(), |
105 | 32 | vector_search: advanced::VectorSearchConfig::default(), |
106 | 32 | cdc: advanced::CdcConfig::default(), |
107 | 32 | ttl: advanced::TtlConfig::default(), |
108 | 32 | tenant_affinity: advanced::TenantAffinityConfig::default(), |
109 | 32 | shadow: advanced::ShadowConfig::default(), |
110 | 32 | ilm: advanced::IlmConfig::default(), |
111 | 32 | canary_runner: advanced::CanaryRunnerConfig::default(), |
112 | 32 | explain: advanced::ExplainConfig::default(), |
113 | 32 | admin_ui: advanced::AdminUiConfig::default(), |
114 | 32 | search_ui: advanced::SearchUiConfig::default(), |
115 | 32 | peer_discovery: PeerDiscoveryConfig::default(), |
116 | 32 | leader_election: LeaderElectionConfig::default(), |
117 | 32 | hpa: HpaConfig::default(), |
118 | 32 | } |
119 | 32 | } |
120 | } | |
121 | ||
122 | impl MiroirConfig { | |
123 | /// Validate cross-field constraints. Returns `Ok(())` or a `ConfigError`. | |
124 | 22 | pub fn validate(&self) -> Result<(), ConfigError> { |
125 | 22 | validate::validate(self) |
126 | 22 | } |
127 | ||
128 | /// Layered loading: file → env overrides → CLI overrides. | |
129 | 0 | pub fn load() -> Result<Self, ConfigError> { |
130 | 0 | load::load() |
131 | 0 | } |
132 | ||
133 | /// Load from a specific file path with env-var overrides applied. | |
134 | 0 | pub fn load_from(path: &std::path::Path) -> Result<Self, ConfigError> { |
135 | 0 | load::load_from(path) |
136 | 0 | } |
137 | ||
138 | /// Load from a YAML string (useful for testing). | |
139 | 0 | pub fn from_yaml(yaml: &str) -> Result<Self, ConfigError> { |
140 | 0 | load::from_yaml(yaml) |
141 | 0 | } |
142 | } | |
143 | ||
144 | // --------------------------------------------------------------------------- | |
145 | // Core sub-structs (§4) | |
146 | // --------------------------------------------------------------------------- | |
147 | ||
148 | /// A single Meilisearch node in the cluster topology. | |
149 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
150 | pub struct NodeConfig { | |
151 | pub id: String, | |
152 | pub address: String, | |
153 | pub replica_group: u32, | |
154 | } | |
155 | ||
156 | /// Task store backend configuration. | |
157 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
158 | #[serde(default)] | |
159 | pub struct TaskStoreConfig { | |
160 | /// `sqlite` or `redis`. | |
161 | pub backend: String, | |
162 | /// Path to SQLite database file (sqlite backend). | |
163 | pub path: String, | |
164 | /// Redis URL (redis backend), e.g. `redis://host:6379`. | |
165 | pub url: String, | |
166 | } | |
167 | ||
168 | impl Default for TaskStoreConfig { | |
169 | 50 | fn default() -> Self { |
170 | 50 | Self { |
171 | 50 | backend: "sqlite".into(), |
172 | 50 | path: "/data/miroir-tasks.db".into(), |
173 | 50 | url: String::new(), |
174 | 50 | } |
175 | 50 | } |
176 | } | |
177 | ||
178 | /// Admin API configuration. | |
179 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
180 | #[serde(default)] | |
181 | pub struct AdminConfig { | |
182 | pub enabled: bool, | |
183 | /// Env override: `MIROIR_ADMIN_API_KEY`. | |
184 | pub api_key: String, | |
185 | } | |
186 | ||
187 | impl Default for AdminConfig { | |
188 | 38 | fn default() -> Self { |
189 | 38 | Self { |
190 | 38 | enabled: true, |
191 | 38 | api_key: String::new(), |
192 | 38 | } |
193 | 38 | } |
194 | } | |
195 | ||
196 | /// Health check configuration. | |
197 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
198 | #[serde(default)] | |
199 | pub struct HealthConfig { | |
200 | pub interval_ms: u64, | |
201 | pub timeout_ms: u64, | |
202 | pub unhealthy_threshold: u32, | |
203 | pub recovery_threshold: u32, | |
204 | } | |
205 | ||
206 | impl Default for HealthConfig { | |
207 | 38 | fn default() -> Self { |
208 | 38 | Self { |
209 | 38 | interval_ms: 5000, |
210 | 38 | timeout_ms: 2000, |
211 | 38 | unhealthy_threshold: 3, |
212 | 38 | recovery_threshold: 2, |
213 | 38 | } |
214 | 38 | } |
215 | } | |
216 | ||
217 | /// Scatter-gather query configuration. | |
218 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
219 | #[serde(default)] | |
220 | pub struct ScatterConfig { | |
221 | pub node_timeout_ms: u64, | |
222 | pub retry_on_timeout: bool, | |
223 | /// `partial` or `error`. | |
224 | pub unavailable_shard_policy: String, | |
225 | } | |
226 | ||
227 | impl Default for ScatterConfig { | |
228 | 38 | fn default() -> Self { |
229 | 38 | Self { |
230 | 38 | node_timeout_ms: 5000, |
231 | 38 | retry_on_timeout: true, |
232 | 38 | unavailable_shard_policy: "partial".into(), |
233 | 38 | } |
234 | 38 | } |
235 | } | |
236 | ||
237 | /// Rebalancer configuration. | |
238 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
239 | #[serde(default)] | |
240 | pub struct RebalancerConfig { | |
241 | pub auto_rebalance_on_recovery: bool, | |
242 | pub max_concurrent_migrations: u32, | |
243 | pub migration_timeout_s: u64, | |
244 | } | |
245 | ||
246 | impl Default for RebalancerConfig { | |
247 | 38 | fn default() -> Self { |
248 | 38 | Self { |
249 | 38 | auto_rebalance_on_recovery: true, |
250 | 38 | max_concurrent_migrations: 4, |
251 | 38 | migration_timeout_s: 3600, |
252 | 38 | } |
253 | 38 | } |
254 | } | |
255 | ||
256 | /// Server (HTTP listener) configuration. | |
257 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
258 | #[serde(default)] | |
259 | pub struct ServerConfig { | |
260 | pub port: u16, | |
261 | pub bind: String, | |
262 | pub max_body_bytes: u64, | |
263 | #[serde(default = "default_max_concurrent_requests")] | |
264 | pub max_concurrent_requests: u32, | |
265 | #[serde(default = "default_request_timeout_ms")] | |
266 | pub request_timeout_ms: u64, | |
267 | } | |
268 | ||
269 | 42 | fn default_max_concurrent_requests() -> u32 { |
270 | 42 | 500 |
271 | 42 | } |
272 | 42 | fn default_request_timeout_ms() -> u64 { |
273 | 42 | 30000 |
274 | 42 | } |
275 | ||
276 | impl Default for ServerConfig { | |
277 | 38 | fn default() -> Self { |
278 | 38 | Self { |
279 | 38 | port: 7700, |
280 | 38 | bind: "0.0.0.0".into(), |
281 | 38 | max_body_bytes: 104_857_600, // 100 MiB |
282 | 38 | max_concurrent_requests: default_max_concurrent_requests(), |
283 | 38 | request_timeout_ms: default_request_timeout_ms(), |
284 | 38 | } |
285 | 38 | } |
286 | } | |
287 | ||
288 | /// HTTP/2 connection pool per-node settings (§14.8). | |
289 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
290 | #[serde(default)] | |
291 | pub struct ConnectionPoolConfig { | |
292 | pub max_idle: u32, | |
293 | pub max_total: u32, | |
294 | pub idle_timeout_s: u64, | |
295 | } | |
296 | ||
297 | impl Default for ConnectionPoolConfig { | |
298 | 34 | fn default() -> Self { |
299 | 34 | Self { |
300 | 34 | max_idle: 32, |
301 | 34 | max_total: 128, |
302 | 34 | idle_timeout_s: 60, |
303 | 34 | } |
304 | 34 | } |
305 | } | |
306 | ||
307 | /// Task registry cache settings (§14.8). | |
308 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
309 | #[serde(default)] | |
310 | pub struct TaskRegistryConfig { | |
311 | pub cache_size: u32, | |
312 | pub redis_pool_max: u32, | |
313 | } | |
314 | ||
315 | impl Default for TaskRegistryConfig { | |
316 | 34 | fn default() -> Self { |
317 | 34 | Self { |
318 | 34 | cache_size: 10000, |
319 | 34 | redis_pool_max: 50, |
320 | 34 | } |
321 | 34 | } |
322 | } | |
323 | ||
324 | /// Peer discovery via Kubernetes headless Service (§14.5). | |
325 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
326 | #[serde(default)] | |
327 | pub struct PeerDiscoveryConfig { | |
328 | pub service_name: String, | |
329 | pub refresh_interval_s: u64, | |
330 | } | |
331 | ||
332 | impl Default for PeerDiscoveryConfig { | |
333 | 34 | fn default() -> Self { |
334 | 34 | Self { |
335 | 34 | service_name: "miroir-headless".into(), |
336 | 34 | refresh_interval_s: 15, |
337 | 34 | } |
338 | 34 | } |
339 | } | |
340 | ||
341 | /// Leader election for Mode B background jobs (§14.5). | |
342 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
343 | #[serde(default)] | |
344 | pub struct LeaderElectionConfig { | |
345 | pub enabled: bool, | |
346 | pub lease_ttl_s: u64, | |
347 | pub renew_interval_s: u64, | |
348 | } | |
349 | ||
350 | impl Default for LeaderElectionConfig { | |
351 | 38 | fn default() -> Self { |
352 | 38 | Self { |
353 | 38 | enabled: true, |
354 | 38 | lease_ttl_s: 10, |
355 | 38 | renew_interval_s: 3, |
356 | 38 | } |
357 | 38 | } |
358 | } | |
359 | ||
360 | /// Horizontal Pod Autoscaler settings (Helm-only, informational in config). | |
361 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] | |
362 | #[serde(default)] | |
363 | pub struct HpaConfig { | |
364 | #[serde(default)] | |
365 | pub enabled: bool, | |
366 | } | |
367 | ||
368 | /// Policy for handling unavailable shards during scatter. | |
369 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] | |
370 | #[serde(rename_all = "snake_case")] | |
371 | pub enum UnavailableShardPolicy { | |
372 | /// Return partial results from available nodes. | |
373 | Partial, | |
374 | /// Fail the request if any shard is unavailable. | |
375 | Error, | |
376 | /// Fall back to another replica group for unavailable shards. | |
377 | Fallback, | |
378 | } | |
379 | ||
380 | impl Default for UnavailableShardPolicy { | |
381 | 0 | fn default() -> Self { |
382 | 0 | Self::Partial |
383 | 0 | } |
384 | } | |
385 | ||
386 | impl std::fmt::Display for UnavailableShardPolicy { | |
387 | 0 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
388 | 0 | match self { |
389 | 0 | Self::Partial => write!(f, "partial"), |
390 | 0 | Self::Error => write!(f, "error"), |
391 | 0 | Self::Fallback => write!(f, "fallback"), |
392 | } | |
393 | 0 | } |
394 | } | |
395 | ||
396 | #[cfg(test)] | |
397 | mod tests { | |
398 | use super::*; | |
399 | ||
400 | /// Returns a minimal valid dev config (single-node, sqlite, RF=1). | |
401 | 10 | fn dev_config() -> MiroirConfig { |
402 | 10 | MiroirConfig { |
403 | 10 | replication_factor: 1, |
404 | 10 | task_store: TaskStoreConfig { |
405 | 10 | backend: "sqlite".into(), |
406 | 10 | ..Default::default() |
407 | 10 | }, |
408 | 10 | cdc: advanced::CdcConfig { |
409 | 10 | buffer: advanced::CdcBufferConfig { |
410 | 10 | overflow: "drop".into(), |
411 | 10 | ..Default::default() |
412 | 10 | }, |
413 | 10 | ..Default::default() |
414 | 10 | }, |
415 | 10 | search_ui: advanced::SearchUiConfig { |
416 | 10 | rate_limit: advanced::SearchUiRateLimitConfig { |
417 | 10 | backend: "local".into(), |
418 | 10 | ..Default::default() |
419 | 10 | }, |
420 | 10 | ..Default::default() |
421 | 10 | }, |
422 | 10 | ..Default::default() |
423 | 10 | } |
424 | 10 | } |
425 | ||
426 | #[test] | |
427 | 2 | fn default_config_is_valid() { |
428 | 2 | let cfg = MiroirConfig::default(); |
429 | // Default has replication_factor=2 with sqlite, which should fail | |
430 | // validation — but the struct itself should construct fine. | |
431 | 2 | assert_eq!(cfg.shards, 64); |
432 | 2 | assert_eq!(cfg.replication_factor, 2); |
433 | 2 | assert_eq!(cfg.replica_groups, 1); |
434 | 2 | assert_eq!(cfg.task_store.backend, "sqlite"); |
435 | 2 | } |
436 | ||
437 | #[test] | |
438 | 2 | fn minimal_yaml_deserializes() { |
439 | 2 | let yaml = r#" |
440 | 2 | shards: 32 |
441 | 2 | replication_factor: 1 |
442 | 2 | nodes: [] |
443 | 2 | "#; |
444 | 2 | let cfg: MiroirConfig = serde_yaml::from_str(yaml).expect("deserialize"); |
445 | 2 | assert_eq!(cfg.shards, 32); |
446 | 2 | assert_eq!(cfg.replication_factor, 1); |
447 | // All §13 blocks should get defaults | |
448 | 2 | assert!(cfg.resharding.enabled); |
449 | 2 | assert!(cfg.hedging.enabled); |
450 | 2 | assert!(cfg.anti_entropy.enabled); |
451 | 2 | } |
452 | ||
453 | #[test] | |
454 | 2 | fn full_plan_example_deserializes() { |
455 | 2 | let yaml = r#" |
456 | 2 | master_key: "test-key" |
457 | 2 | node_master_key: "node-key" |
458 | 2 | shards: 64 |
459 | 2 | replication_factor: 2 |
460 | 2 | replica_groups: 2 |
461 | 2 | task_store: |
462 | 2 | backend: redis |
463 | 2 | url: "redis://redis:6379" |
464 | 2 | admin: |
465 | 2 | enabled: true |
466 | 2 | nodes: |
467 | 2 | - id: "meili-0" |
468 | 2 | address: "http://meili-0.search.svc:7700" |
469 | 2 | replica_group: 0 |
470 | 2 | - id: "meili-1" |
471 | 2 | address: "http://meili-1.search.svc:7700" |
472 | 2 | replica_group: 0 |
473 | 2 | health: |
474 | 2 | interval_ms: 5000 |
475 | 2 | timeout_ms: 2000 |
476 | 2 | unhealthy_threshold: 3 |
477 | 2 | recovery_threshold: 2 |
478 | 2 | scatter: |
479 | 2 | node_timeout_ms: 5000 |
480 | 2 | retry_on_timeout: true |
481 | 2 | unavailable_shard_policy: partial |
482 | 2 | rebalancer: |
483 | 2 | auto_rebalance_on_recovery: true |
484 | 2 | max_concurrent_migrations: 4 |
485 | 2 | migration_timeout_s: 3600 |
486 | 2 | server: |
487 | 2 | port: 7700 |
488 | 2 | bind: "0.0.0.0" |
489 | 2 | max_body_bytes: 104857600 |
490 | 2 | leader_election: |
491 | 2 | enabled: true |
492 | 2 | "#; |
493 | 2 | let cfg: MiroirConfig = serde_yaml::from_str(yaml).expect("deserialize"); |
494 | 2 | assert_eq!(cfg.master_key, "test-key"); |
495 | 2 | assert_eq!(cfg.nodes.len(), 2); |
496 | 2 | assert_eq!(cfg.replica_groups, 2); |
497 | 2 | cfg.validate().expect("valid production config"); |
498 | 2 | } |
499 | ||
500 | #[test] | |
501 | 2 | fn round_trip_yaml() { |
502 | 2 | let original = MiroirConfig::default(); |
503 | 2 | let yaml = serde_yaml::to_string(&original).expect("serialize"); |
504 | 2 | let round_tripped: MiroirConfig = serde_yaml::from_str(&yaml).expect("deserialize"); |
505 | 2 | assert_eq!(original, round_tripped); |
506 | 2 | } |
507 | ||
508 | #[test] | |
509 | 2 | fn validation_rejects_ha_with_sqlite() { |
510 | 2 | let mut cfg = dev_config(); |
511 | 2 | cfg.replication_factor = 2; |
512 | 2 | let err = cfg.validate().unwrap_err(); |
513 | 2 | assert!(err.to_string().contains("redis")); |
514 | 2 | } |
515 | ||
516 | #[test] | |
517 | 2 | fn validation_rejects_zero_shards() { |
518 | 2 | let mut cfg = dev_config(); |
519 | 2 | cfg.shards = 0; |
520 | 2 | let err = cfg.validate().unwrap_err(); |
521 | 2 | assert!(err.to_string().contains("shards")); |
522 | 2 | } |
523 | ||
524 | #[test] | |
525 | 2 | fn validation_rejects_duplicate_node_ids() { |
526 | 2 | let mut cfg = dev_config(); |
527 | 2 | cfg.nodes = vec![ |
528 | 2 | NodeConfig { |
529 | 2 | id: "n0".into(), |
530 | 2 | address: "http://n0".into(), |
531 | 2 | replica_group: 0, |
532 | 2 | }, |
533 | 2 | NodeConfig { |
534 | 2 | id: "n0".into(), |
535 | 2 | address: "http://n0b".into(), |
536 | 2 | replica_group: 0, |
537 | 2 | }, |
538 | ]; | |
539 | 2 | let err = cfg.validate().unwrap_err(); |
540 | 2 | assert!(err.to_string().contains("duplicate")); |
541 | 2 | } |
542 | ||
543 | #[test] | |
544 | 2 | fn validation_rejects_node_outside_replica_groups() { |
545 | 2 | let mut cfg = dev_config(); |
546 | 2 | cfg.nodes = vec![NodeConfig { |
547 | 2 | id: "n0".into(), |
548 | 2 | address: "http://n0".into(), |
549 | 2 | replica_group: 5, |
550 | 2 | }]; |
551 | 2 | let err = cfg.validate().unwrap_err(); |
552 | 2 | assert!(err.to_string().contains("replica_group")); |
553 | 2 | } |
554 | ||
555 | #[test] | |
556 | 2 | fn validation_rejects_scoped_key_timing_inversion() { |
557 | 2 | let mut cfg = dev_config(); |
558 | 2 | cfg.search_ui.scoped_key_max_age_days = 10; |
559 | 2 | cfg.search_ui.scoped_key_rotate_before_expiry_days = 10; |
560 | 2 | let err = cfg.validate().unwrap_err(); |
561 | 2 | assert!(err.to_string().contains("scoped_key")); |
562 | 2 | } |
563 | ||
564 | #[test] | |
565 | 2 | fn advanced_defaults_all_enabled() { |
566 | 2 | let cfg = MiroirConfig::default(); |
567 | 2 | assert!(cfg.resharding.enabled); |
568 | 2 | assert!(cfg.hedging.enabled); |
569 | 2 | assert!(cfg.replica_selection.strategy == "adaptive"); |
570 | 2 | assert!(cfg.query_planner.enabled); |
571 | 2 | assert!(cfg.settings_broadcast.strategy == "two_phase"); |
572 | 2 | assert!(cfg.session_pinning.enabled); |
573 | 2 | assert!(cfg.aliases.enabled); |
574 | 2 | assert!(cfg.anti_entropy.enabled); |
575 | 2 | assert!(cfg.dump_import.mode == "streaming"); |
576 | 2 | assert!(cfg.idempotency.enabled); |
577 | 2 | assert!(cfg.query_coalescing.enabled); |
578 | 2 | assert!(cfg.multi_search.enabled); |
579 | 2 | assert!(cfg.vector_search.enabled); |
580 | 2 | assert!(cfg.cdc.enabled); |
581 | 2 | assert!(cfg.ttl.enabled); |
582 | 2 | assert!(cfg.tenant_affinity.enabled); |
583 | 2 | assert!(cfg.shadow.enabled); |
584 | 2 | assert!(cfg.ilm.enabled); |
585 | 2 | assert!(cfg.canary_runner.enabled); |
586 | 2 | assert!(cfg.explain.enabled); |
587 | 2 | assert!(cfg.admin_ui.enabled); |
588 | 2 | assert!(cfg.search_ui.enabled); |
589 | 2 | } |
590 | } |
Line | Count | Source |
1 | //! §13 Advanced capabilities configuration structs. | |
2 | ||
3 | use serde::{Deserialize, Serialize}; | |
4 | use std::collections::HashMap; | |
5 | ||
6 | // --------------------------------------------------------------------------- | |
7 | // 13.1 Online resharding | |
8 | // --------------------------------------------------------------------------- | |
9 | ||
10 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
11 | #[serde(default)] | |
12 | pub struct ReshardingConfig { | |
13 | pub enabled: bool, | |
14 | pub backfill_concurrency: u32, | |
15 | pub backfill_batch_size: u32, | |
16 | pub throttle_docs_per_sec: u32, | |
17 | pub verify_before_swap: bool, | |
18 | pub retain_old_index_hours: u32, | |
19 | } | |
20 | ||
21 | impl Default for ReshardingConfig { | |
22 | 34 | fn default() -> Self { |
23 | 34 | Self { |
24 | 34 | enabled: true, |
25 | 34 | backfill_concurrency: 4, |
26 | 34 | backfill_batch_size: 1000, |
27 | 34 | throttle_docs_per_sec: 0, |
28 | 34 | verify_before_swap: true, |
29 | 34 | retain_old_index_hours: 48, |
30 | 34 | } |
31 | 34 | } |
32 | } | |
33 | ||
34 | // --------------------------------------------------------------------------- | |
35 | // 13.2 Hedged requests | |
36 | // --------------------------------------------------------------------------- | |
37 | ||
38 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
39 | #[serde(default)] | |
40 | pub struct HedgingConfig { | |
41 | pub enabled: bool, | |
42 | pub p95_trigger_multiplier: f64, | |
43 | pub min_trigger_ms: u64, | |
44 | pub max_hedges_per_query: u32, | |
45 | pub cross_group_fallback: bool, | |
46 | } | |
47 | ||
48 | impl Default for HedgingConfig { | |
49 | 34 | fn default() -> Self { |
50 | 34 | Self { |
51 | 34 | enabled: true, |
52 | 34 | p95_trigger_multiplier: 1.2, |
53 | 34 | min_trigger_ms: 15, |
54 | 34 | max_hedges_per_query: 2, |
55 | 34 | cross_group_fallback: true, |
56 | 34 | } |
57 | 34 | } |
58 | } | |
59 | ||
60 | // --------------------------------------------------------------------------- | |
61 | // 13.3 Adaptive replica selection (EWMA) | |
62 | // --------------------------------------------------------------------------- | |
63 | ||
64 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
65 | #[serde(default)] | |
66 | pub struct ReplicaSelectionConfig { | |
67 | /// `adaptive`, `round_robin`, or `random`. | |
68 | pub strategy: String, | |
69 | pub latency_weight: f64, | |
70 | pub inflight_weight: f64, | |
71 | pub error_weight: f64, | |
72 | pub ewma_half_life_ms: u64, | |
73 | pub exploration_epsilon: f64, | |
74 | } | |
75 | ||
76 | impl Default for ReplicaSelectionConfig { | |
77 | 34 | fn default() -> Self { |
78 | 34 | Self { |
79 | 34 | strategy: "adaptive".into(), |
80 | 34 | latency_weight: 1.0, |
81 | 34 | inflight_weight: 2.0, |
82 | 34 | error_weight: 10.0, |
83 | 34 | ewma_half_life_ms: 5000, |
84 | 34 | exploration_epsilon: 0.05, |
85 | 34 | } |
86 | 34 | } |
87 | } | |
88 | ||
89 | // --------------------------------------------------------------------------- | |
90 | // 13.4 Shard-aware query planner | |
91 | // --------------------------------------------------------------------------- | |
92 | ||
93 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
94 | #[serde(default)] | |
95 | pub struct QueryPlannerConfig { | |
96 | pub enabled: bool, | |
97 | pub max_pk_literals_narrowable: u32, | |
98 | pub log_plans: bool, | |
99 | } | |
100 | ||
101 | impl Default for QueryPlannerConfig { | |
102 | 34 | fn default() -> Self { |
103 | 34 | Self { |
104 | 34 | enabled: true, |
105 | 34 | max_pk_literals_narrowable: 128, |
106 | 34 | log_plans: false, |
107 | 34 | } |
108 | 34 | } |
109 | } | |
110 | ||
111 | // --------------------------------------------------------------------------- | |
112 | // 13.5 Two-phase settings broadcast | |
113 | // --------------------------------------------------------------------------- | |
114 | ||
115 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
116 | #[serde(default)] | |
117 | pub struct SettingsBroadcastConfig { | |
118 | /// `two_phase` or `sequential` (legacy). | |
119 | pub strategy: String, | |
120 | pub verify_timeout_s: u64, | |
121 | pub max_repair_retries: u32, | |
122 | pub freeze_writes_on_unrepairable: bool, | |
123 | } | |
124 | ||
125 | impl Default for SettingsBroadcastConfig { | |
126 | 34 | fn default() -> Self { |
127 | 34 | Self { |
128 | 34 | strategy: "two_phase".into(), |
129 | 34 | verify_timeout_s: 60, |
130 | 34 | max_repair_retries: 3, |
131 | 34 | freeze_writes_on_unrepairable: true, |
132 | 34 | } |
133 | 34 | } |
134 | } | |
135 | ||
136 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
137 | #[serde(default)] | |
138 | pub struct SettingsDriftCheckConfig { | |
139 | pub interval_s: u64, | |
140 | pub auto_repair: bool, | |
141 | } | |
142 | ||
143 | impl Default for SettingsDriftCheckConfig { | |
144 | 34 | fn default() -> Self { |
145 | 34 | Self { |
146 | 34 | interval_s: 300, |
147 | 34 | auto_repair: true, |
148 | 34 | } |
149 | 34 | } |
150 | } | |
151 | ||
152 | // --------------------------------------------------------------------------- | |
153 | // 13.6 Session pinning (read-your-writes) | |
154 | // --------------------------------------------------------------------------- | |
155 | ||
156 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
157 | #[serde(default)] | |
158 | pub struct SessionPinningConfig { | |
159 | pub enabled: bool, | |
160 | pub ttl_seconds: u64, | |
161 | pub max_sessions: u32, | |
162 | /// `block` or `route_pin`. | |
163 | pub wait_strategy: String, | |
164 | pub max_wait_ms: u64, | |
165 | } | |
166 | ||
167 | impl Default for SessionPinningConfig { | |
168 | 34 | fn default() -> Self { |
169 | 34 | Self { |
170 | 34 | enabled: true, |
171 | 34 | ttl_seconds: 900, |
172 | 34 | max_sessions: 100_000, |
173 | 34 | wait_strategy: "block".into(), |
174 | 34 | max_wait_ms: 5000, |
175 | 34 | } |
176 | 34 | } |
177 | } | |
178 | ||
179 | // --------------------------------------------------------------------------- | |
180 | // 13.7 Index aliases | |
181 | // --------------------------------------------------------------------------- | |
182 | ||
183 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
184 | #[serde(default)] | |
185 | pub struct AliasesConfig { | |
186 | pub enabled: bool, | |
187 | pub history_retention: u32, | |
188 | pub require_target_exists: bool, | |
189 | } | |
190 | ||
191 | impl Default for AliasesConfig { | |
192 | 34 | fn default() -> Self { |
193 | 34 | Self { |
194 | 34 | enabled: true, |
195 | 34 | history_retention: 10, |
196 | 34 | require_target_exists: true, |
197 | 34 | } |
198 | 34 | } |
199 | } | |
200 | ||
201 | // --------------------------------------------------------------------------- | |
202 | // 13.8 Anti-entropy shard reconciler | |
203 | // --------------------------------------------------------------------------- | |
204 | ||
205 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
206 | #[serde(default)] | |
207 | pub struct AntiEntropyConfig { | |
208 | pub enabled: bool, | |
209 | pub schedule: String, | |
210 | pub shards_per_pass: u32, | |
211 | pub max_read_concurrency: u32, | |
212 | pub fingerprint_batch_size: u32, | |
213 | pub auto_repair: bool, | |
214 | pub updated_at_field: String, | |
215 | } | |
216 | ||
217 | impl Default for AntiEntropyConfig { | |
218 | 34 | fn default() -> Self { |
219 | 34 | Self { |
220 | 34 | enabled: true, |
221 | 34 | schedule: "every 6h".into(), |
222 | 34 | shards_per_pass: 0, |
223 | 34 | max_read_concurrency: 2, |
224 | 34 | fingerprint_batch_size: 1000, |
225 | 34 | auto_repair: true, |
226 | 34 | updated_at_field: "_miroir_updated_at".into(), |
227 | 34 | } |
228 | 34 | } |
229 | } | |
230 | ||
231 | // --------------------------------------------------------------------------- | |
232 | // 13.9 Streaming dump import | |
233 | // --------------------------------------------------------------------------- | |
234 | ||
235 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
236 | #[serde(default)] | |
237 | pub struct DumpImportConfig { | |
238 | /// `streaming` or `broadcast` (legacy). | |
239 | pub mode: String, | |
240 | pub batch_size: u32, | |
241 | pub parallel_target_writes: u32, | |
242 | pub memory_buffer_bytes: u64, | |
243 | pub chunk_size_bytes: u64, | |
244 | } | |
245 | ||
246 | impl Default for DumpImportConfig { | |
247 | 34 | fn default() -> Self { |
248 | 34 | Self { |
249 | 34 | mode: "streaming".into(), |
250 | 34 | batch_size: 1000, |
251 | 34 | parallel_target_writes: 8, |
252 | 34 | memory_buffer_bytes: 134_217_728, // 128 MiB |
253 | 34 | chunk_size_bytes: 268_435_456, // 256 MiB |
254 | 34 | } |
255 | 34 | } |
256 | } | |
257 | ||
258 | // --------------------------------------------------------------------------- | |
259 | // 13.10 Idempotency keys | |
260 | // --------------------------------------------------------------------------- | |
261 | ||
262 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
263 | #[serde(default)] | |
264 | pub struct IdempotencyConfig { | |
265 | pub enabled: bool, | |
266 | pub ttl_seconds: u64, | |
267 | pub max_cached_keys: u32, | |
268 | } | |
269 | ||
270 | impl Default for IdempotencyConfig { | |
271 | 34 | fn default() -> Self { |
272 | 34 | Self { |
273 | 34 | enabled: true, |
274 | 34 | ttl_seconds: 86400, |
275 | 34 | max_cached_keys: 1_000_000, |
276 | 34 | } |
277 | 34 | } |
278 | } | |
279 | ||
280 | // --------------------------------------------------------------------------- | |
281 | // 13.10 Query coalescing (paired with idempotency) | |
282 | // --------------------------------------------------------------------------- | |
283 | ||
284 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
285 | #[serde(default)] | |
286 | pub struct QueryCoalescingConfig { | |
287 | pub enabled: bool, | |
288 | pub window_ms: u64, | |
289 | pub max_subscribers: u32, | |
290 | pub max_pending_queries: u32, | |
291 | } | |
292 | ||
293 | impl Default for QueryCoalescingConfig { | |
294 | 34 | fn default() -> Self { |
295 | 34 | Self { |
296 | 34 | enabled: true, |
297 | 34 | window_ms: 50, |
298 | 34 | max_subscribers: 1000, |
299 | 34 | max_pending_queries: 10000, |
300 | 34 | } |
301 | 34 | } |
302 | } | |
303 | ||
304 | // --------------------------------------------------------------------------- | |
305 | // 13.11 Multi-search batch API | |
306 | // --------------------------------------------------------------------------- | |
307 | ||
308 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
309 | #[serde(default)] | |
310 | pub struct MultiSearchConfig { | |
311 | pub enabled: bool, | |
312 | pub max_queries_per_batch: u32, | |
313 | pub total_timeout_ms: u64, | |
314 | pub per_query_timeout_ms: u64, | |
315 | } | |
316 | ||
317 | impl Default for MultiSearchConfig { | |
318 | 34 | fn default() -> Self { |
319 | 34 | Self { |
320 | 34 | enabled: true, |
321 | 34 | max_queries_per_batch: 100, |
322 | 34 | total_timeout_ms: 30000, |
323 | 34 | per_query_timeout_ms: 30000, |
324 | 34 | } |
325 | 34 | } |
326 | } | |
327 | ||
328 | // --------------------------------------------------------------------------- | |
329 | // 13.12 Vector / hybrid search | |
330 | // --------------------------------------------------------------------------- | |
331 | ||
332 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
333 | #[serde(default)] | |
334 | pub struct VectorSearchConfig { | |
335 | pub enabled: bool, | |
336 | pub over_fetch_factor: u32, | |
337 | /// `convex` or `rrf`. | |
338 | pub merge_strategy: String, | |
339 | pub hybrid_alpha_default: f64, | |
340 | pub rrf_k: u32, | |
341 | } | |
342 | ||
343 | impl Default for VectorSearchConfig { | |
344 | 34 | fn default() -> Self { |
345 | 34 | Self { |
346 | 34 | enabled: true, |
347 | 34 | over_fetch_factor: 3, |
348 | 34 | merge_strategy: "convex".into(), |
349 | 34 | hybrid_alpha_default: 0.5, |
350 | 34 | rrf_k: 60, |
351 | 34 | } |
352 | 34 | } |
353 | } | |
354 | ||
355 | // --------------------------------------------------------------------------- | |
356 | // 13.13 Change data capture (CDC) | |
357 | // --------------------------------------------------------------------------- | |
358 | ||
359 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
360 | #[serde(default)] | |
361 | pub struct CdcConfig { | |
362 | pub enabled: bool, | |
363 | pub emit_ttl_deletes: bool, | |
364 | pub emit_internal_writes: bool, | |
365 | pub sinks: Vec<CdcSinkConfig>, | |
366 | pub buffer: CdcBufferConfig, | |
367 | } | |
368 | ||
369 | impl Default for CdcConfig { | |
370 | 46 | fn default() -> Self { |
371 | 46 | Self { |
372 | 46 | enabled: true, |
373 | 46 | emit_ttl_deletes: false, |
374 | 46 | emit_internal_writes: false, |
375 | 46 | sinks: Vec::new(), |
376 | 46 | buffer: CdcBufferConfig::default(), |
377 | 46 | } |
378 | 46 | } |
379 | } | |
380 | ||
381 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
382 | #[serde(default)] | |
383 | pub struct CdcSinkConfig { | |
384 | /// `webhook`, `nats`, `kafka`, or `internal`. | |
385 | #[serde(rename = "type")] | |
386 | pub sink_type: String, | |
387 | pub url: String, | |
388 | pub batch_size: u32, | |
389 | pub batch_flush_ms: u64, | |
390 | pub include_body: bool, | |
391 | pub retry_max_s: u64, | |
392 | /// NATS-specific. | |
393 | pub subject_prefix: Option<String>, | |
394 | } | |
395 | ||
396 | impl Default for CdcSinkConfig { | |
397 | 0 | fn default() -> Self { |
398 | 0 | Self { |
399 | 0 | sink_type: "webhook".into(), |
400 | 0 | url: String::new(), |
401 | 0 | batch_size: 100, |
402 | 0 | batch_flush_ms: 1000, |
403 | 0 | include_body: false, |
404 | 0 | retry_max_s: 3600, |
405 | 0 | subject_prefix: None, |
406 | 0 | } |
407 | 0 | } |
408 | } | |
409 | ||
410 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
411 | #[serde(default)] | |
412 | pub struct CdcBufferConfig { | |
413 | /// `memory`, `redis`, or `pvc`. | |
414 | pub primary: String, | |
415 | pub memory_bytes: u64, | |
416 | /// `redis`, `pvc`, or `drop`. | |
417 | pub overflow: String, | |
418 | pub redis_bytes: u64, | |
419 | } | |
420 | ||
421 | impl Default for CdcBufferConfig { | |
422 | 58 | fn default() -> Self { |
423 | 58 | Self { |
424 | 58 | primary: "memory".into(), |
425 | 58 | memory_bytes: 67_108_864, // 64 MiB |
426 | 58 | overflow: "redis".into(), |
427 | 58 | redis_bytes: 1_073_741_824, // 1 GiB |
428 | 58 | } |
429 | 58 | } |
430 | } | |
431 | ||
432 | // --------------------------------------------------------------------------- | |
433 | // 13.14 Document TTL | |
434 | // --------------------------------------------------------------------------- | |
435 | ||
436 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
437 | #[serde(default)] | |
438 | pub struct TtlConfig { | |
439 | pub enabled: bool, | |
440 | pub sweep_interval_s: u64, | |
441 | pub max_deletes_per_sweep: u32, | |
442 | pub expires_at_field: String, | |
443 | pub per_index_overrides: HashMap<String, TtlOverride>, | |
444 | } | |
445 | ||
446 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
447 | pub struct TtlOverride { | |
448 | pub sweep_interval_s: u64, | |
449 | pub max_deletes_per_sweep: u32, | |
450 | } | |
451 | ||
452 | impl Default for TtlConfig { | |
453 | 34 | fn default() -> Self { |
454 | 34 | Self { |
455 | 34 | enabled: true, |
456 | 34 | sweep_interval_s: 300, |
457 | 34 | max_deletes_per_sweep: 10000, |
458 | 34 | expires_at_field: "_miroir_expires_at".into(), |
459 | 34 | per_index_overrides: HashMap::new(), |
460 | 34 | } |
461 | 34 | } |
462 | } | |
463 | ||
464 | // --------------------------------------------------------------------------- | |
465 | // 13.15 Tenant-to-replica-group affinity | |
466 | // --------------------------------------------------------------------------- | |
467 | ||
468 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
469 | #[serde(default)] | |
470 | pub struct TenantAffinityConfig { | |
471 | pub enabled: bool, | |
472 | /// `header`, `api_key`, or `explicit`. | |
473 | pub mode: String, | |
474 | pub header_name: String, | |
475 | /// `hash`, `random`, or `reject`. | |
476 | pub fallback: String, | |
477 | pub static_map: HashMap<String, u32>, | |
478 | pub dedicated_groups: Vec<u32>, | |
479 | } | |
480 | ||
481 | impl Default for TenantAffinityConfig { | |
482 | 34 | fn default() -> Self { |
483 | 34 | Self { |
484 | 34 | enabled: true, |
485 | 34 | mode: "header".into(), |
486 | 34 | header_name: "X-Miroir-Tenant".into(), |
487 | 34 | fallback: "hash".into(), |
488 | 34 | static_map: HashMap::new(), |
489 | 34 | dedicated_groups: Vec::new(), |
490 | 34 | } |
491 | 34 | } |
492 | } | |
493 | ||
494 | // --------------------------------------------------------------------------- | |
495 | // 13.16 Traffic shadow | |
496 | // --------------------------------------------------------------------------- | |
497 | ||
498 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
499 | #[serde(default)] | |
500 | pub struct ShadowConfig { | |
501 | pub enabled: bool, | |
502 | pub targets: Vec<ShadowTargetConfig>, | |
503 | pub diff_buffer_size: u32, | |
504 | pub max_shadow_latency_ms: u64, | |
505 | } | |
506 | ||
507 | impl Default for ShadowConfig { | |
508 | 34 | fn default() -> Self { |
509 | 34 | Self { |
510 | 34 | enabled: true, |
511 | 34 | targets: Vec::new(), |
512 | 34 | diff_buffer_size: 10000, |
513 | 34 | max_shadow_latency_ms: 5000, |
514 | 34 | } |
515 | 34 | } |
516 | } | |
517 | ||
518 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
519 | #[serde(default)] | |
520 | pub struct ShadowTargetConfig { | |
521 | pub name: String, | |
522 | pub url: String, | |
523 | pub api_key_env: String, | |
524 | pub sample_rate: f64, | |
525 | pub operations: Vec<String>, | |
526 | } | |
527 | ||
528 | impl Default for ShadowTargetConfig { | |
529 | 0 | fn default() -> Self { |
530 | 0 | Self { |
531 | 0 | name: String::new(), |
532 | 0 | url: String::new(), |
533 | 0 | api_key_env: String::new(), |
534 | 0 | sample_rate: 0.05, |
535 | 0 | operations: vec!["search".into(), "multi_search".into(), "explain".into()], |
536 | 0 | } |
537 | 0 | } |
538 | } | |
539 | ||
540 | // --------------------------------------------------------------------------- | |
541 | // 13.17 Index lifecycle management (ILM) | |
542 | // --------------------------------------------------------------------------- | |
543 | ||
544 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
545 | #[serde(default)] | |
546 | pub struct IlmConfig { | |
547 | pub enabled: bool, | |
548 | pub check_interval_s: u64, | |
549 | pub safety_lock_older_than_days: u32, | |
550 | pub max_rollovers_per_check: u32, | |
551 | } | |
552 | ||
553 | impl Default for IlmConfig { | |
554 | 34 | fn default() -> Self { |
555 | 34 | Self { |
556 | 34 | enabled: true, |
557 | 34 | check_interval_s: 3600, |
558 | 34 | safety_lock_older_than_days: 7, |
559 | 34 | max_rollovers_per_check: 10, |
560 | 34 | } |
561 | 34 | } |
562 | } | |
563 | ||
564 | // --------------------------------------------------------------------------- | |
565 | // 13.18 Synthetic canary queries | |
566 | // --------------------------------------------------------------------------- | |
567 | ||
568 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
569 | #[serde(default)] | |
570 | pub struct CanaryRunnerConfig { | |
571 | pub enabled: bool, | |
572 | pub max_concurrent_canaries: u32, | |
573 | pub run_history_per_canary: u32, | |
574 | pub emit_results_to_cdc: bool, | |
575 | } | |
576 | ||
577 | impl Default for CanaryRunnerConfig { | |
578 | 34 | fn default() -> Self { |
579 | 34 | Self { |
580 | 34 | enabled: true, |
581 | 34 | max_concurrent_canaries: 10, |
582 | 34 | run_history_per_canary: 100, |
583 | 34 | emit_results_to_cdc: true, |
584 | 34 | } |
585 | 34 | } |
586 | } | |
587 | ||
588 | // --------------------------------------------------------------------------- | |
589 | // 13.19 Admin Web UI | |
590 | // --------------------------------------------------------------------------- | |
591 | ||
592 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
593 | #[serde(default)] | |
594 | pub struct AdminUiConfig { | |
595 | pub enabled: bool, | |
596 | pub path: String, | |
597 | /// `key`, `oauth` (future), or `none` (dev only). | |
598 | pub auth: String, | |
599 | pub session_ttl_s: u64, | |
600 | pub read_only_mode: bool, | |
601 | pub allowed_origins: Vec<String>, | |
602 | pub cors_allowed_origins: Vec<String>, | |
603 | pub csp_overrides: CspOverridesConfig, | |
604 | pub theme: AdminUiThemeConfig, | |
605 | pub features: AdminUiFeaturesConfig, | |
606 | } | |
607 | ||
608 | impl Default for AdminUiConfig { | |
609 | 34 | fn default() -> Self { |
610 | 34 | Self { |
611 | 34 | enabled: true, |
612 | 34 | path: "/_miroir/admin".into(), |
613 | 34 | auth: "key".into(), |
614 | 34 | session_ttl_s: 3600, |
615 | 34 | read_only_mode: false, |
616 | 34 | allowed_origins: vec!["same-origin".into()], |
617 | 34 | cors_allowed_origins: Vec::new(), |
618 | 34 | csp_overrides: CspOverridesConfig::default(), |
619 | 34 | theme: AdminUiThemeConfig::default(), |
620 | 34 | features: AdminUiFeaturesConfig::default(), |
621 | 34 | } |
622 | 34 | } |
623 | } | |
624 | ||
625 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] | |
626 | #[serde(default)] | |
627 | pub struct CspOverridesConfig { | |
628 | pub script_src: Vec<String>, | |
629 | pub img_src: Vec<String>, | |
630 | pub connect_src: Vec<String>, | |
631 | } | |
632 | ||
633 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
634 | #[serde(default)] | |
635 | pub struct AdminUiThemeConfig { | |
636 | pub accent_color: String, | |
637 | /// `auto`, `light`, or `dark`. | |
638 | pub default_mode: String, | |
639 | } | |
640 | ||
641 | impl Default for AdminUiThemeConfig { | |
642 | 36 | fn default() -> Self { |
643 | 36 | Self { |
644 | 36 | accent_color: "#2563eb".into(), |
645 | 36 | default_mode: "auto".into(), |
646 | 36 | } |
647 | 36 | } |
648 | } | |
649 | ||
650 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
651 | #[serde(default)] | |
652 | pub struct AdminUiFeaturesConfig { | |
653 | pub sandbox: bool, | |
654 | pub shadow_viewer: bool, | |
655 | pub cdc_inspector: bool, | |
656 | } | |
657 | ||
658 | impl Default for AdminUiFeaturesConfig { | |
659 | 36 | fn default() -> Self { |
660 | 36 | Self { |
661 | 36 | sandbox: true, |
662 | 36 | shadow_viewer: true, |
663 | 36 | cdc_inspector: true, |
664 | 36 | } |
665 | 36 | } |
666 | } | |
667 | ||
668 | // --------------------------------------------------------------------------- | |
669 | // 13.20 Query explain API | |
670 | // --------------------------------------------------------------------------- | |
671 | ||
672 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
673 | #[serde(default)] | |
674 | pub struct ExplainConfig { | |
675 | pub enabled: bool, | |
676 | pub max_warnings: u32, | |
677 | pub allow_execute_parameter: bool, | |
678 | } | |
679 | ||
680 | impl Default for ExplainConfig { | |
681 | 34 | fn default() -> Self { |
682 | 34 | Self { |
683 | 34 | enabled: true, |
684 | 34 | max_warnings: 20, |
685 | 34 | allow_execute_parameter: true, |
686 | 34 | } |
687 | 34 | } |
688 | } | |
689 | ||
690 | // --------------------------------------------------------------------------- | |
691 | // 13.21 Search UI (end-user) | |
692 | // --------------------------------------------------------------------------- | |
693 | ||
694 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
695 | #[serde(default)] | |
696 | pub struct SearchUiConfig { | |
697 | pub enabled: bool, | |
698 | pub path: String, | |
699 | pub widget_script_enabled: bool, | |
700 | pub embeddable: bool, | |
701 | pub auth: SearchUiAuthConfig, | |
702 | pub allowed_origins: Vec<String>, | |
703 | pub scoped_key_max_age_days: u32, | |
704 | pub scoped_key_rotate_before_expiry_days: u32, | |
705 | pub scoped_key_rotation_drain_s: u64, | |
706 | pub rate_limit: SearchUiRateLimitConfig, | |
707 | pub cors_allowed_origins: Vec<String>, | |
708 | pub csp_overrides: CspOverridesConfig, | |
709 | pub csp: String, | |
710 | pub analytics: SearchUiAnalyticsConfig, | |
711 | } | |
712 | ||
713 | impl Default for SearchUiConfig { | |
714 | 46 | fn default() -> Self { |
715 | 46 | Self { |
716 | 46 | enabled: true, |
717 | 46 | path: "/ui/search".into(), |
718 | 46 | widget_script_enabled: true, |
719 | 46 | embeddable: true, |
720 | 46 | auth: SearchUiAuthConfig::default(), |
721 | 46 | allowed_origins: vec!["*".into()], |
722 | 46 | scoped_key_max_age_days: 60, |
723 | 46 | scoped_key_rotate_before_expiry_days: 30, |
724 | 46 | scoped_key_rotation_drain_s: 120, |
725 | 46 | rate_limit: SearchUiRateLimitConfig::default(), |
726 | 46 | cors_allowed_origins: Vec::new(), |
727 | 46 | csp_overrides: CspOverridesConfig::default(), |
728 | 46 | csp: "default-src 'self'; img-src 'self' https:; style-src 'self' 'unsafe-inline'" |
729 | 46 | .into(), |
730 | 46 | analytics: SearchUiAnalyticsConfig::default(), |
731 | 46 | } |
732 | 46 | } |
733 | } | |
734 | ||
735 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
736 | #[serde(default)] | |
737 | pub struct SearchUiAuthConfig { | |
738 | /// `public`, `shared_key`, or `oauth_proxy`. | |
739 | pub mode: String, | |
740 | pub shared_key_env: String, | |
741 | pub session_ttl_s: u64, | |
742 | pub session_rate_limit: String, | |
743 | pub jwt_secret_env: String, | |
744 | pub oauth_proxy: OAuthProxyConfig, | |
745 | } | |
746 | ||
747 | impl Default for SearchUiAuthConfig { | |
748 | 48 | fn default() -> Self { |
749 | 48 | Self { |
750 | 48 | mode: "public".into(), |
751 | 48 | shared_key_env: String::new(), |
752 | 48 | session_ttl_s: 900, |
753 | 48 | session_rate_limit: "10/minute".into(), |
754 | 48 | jwt_secret_env: "SEARCH_UI_JWT_SECRET".into(), |
755 | 48 | oauth_proxy: OAuthProxyConfig::default(), |
756 | 48 | } |
757 | 48 | } |
758 | } | |
759 | ||
760 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
761 | #[serde(default)] | |
762 | pub struct OAuthProxyConfig { | |
763 | pub user_header: String, | |
764 | pub groups_header: String, | |
765 | pub filter_template: Option<String>, | |
766 | pub attribute_map: HashMap<String, String>, | |
767 | } | |
768 | ||
769 | impl Default for OAuthProxyConfig { | |
770 | 50 | fn default() -> Self { |
771 | 50 | Self { |
772 | 50 | user_header: "X-Forwarded-User".into(), |
773 | 50 | groups_header: "X-Forwarded-Groups".into(), |
774 | 50 | filter_template: Some("tenant IN [{groups}]".into()), |
775 | 50 | attribute_map: { |
776 | 50 | let mut m = HashMap::new(); |
777 | 50 | m.insert("groups".into(), "groups_array".into()); |
778 | 50 | m.insert("user".into(), "user_id_string".into()); |
779 | 50 | m |
780 | 50 | }, |
781 | 50 | } |
782 | 50 | } |
783 | } | |
784 | ||
785 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
786 | #[serde(default)] | |
787 | pub struct SearchUiRateLimitConfig { | |
788 | pub per_ip: String, | |
789 | /// `redis` or `local`. | |
790 | pub backend: String, | |
791 | pub redis_key_prefix: String, | |
792 | pub redis_ttl_s: u64, | |
793 | } | |
794 | ||
795 | impl Default for SearchUiRateLimitConfig { | |
796 | 60 | fn default() -> Self { |
797 | 60 | Self { |
798 | 60 | per_ip: "60/minute".into(), |
799 | 60 | backend: "redis".into(), |
800 | 60 | redis_key_prefix: "miroir:ratelimit:searchui:".into(), |
801 | 60 | redis_ttl_s: 60, |
802 | 60 | } |
803 | 60 | } |
804 | } | |
805 | ||
806 | #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] | |
807 | #[serde(default)] | |
808 | pub struct SearchUiAnalyticsConfig { | |
809 | pub enabled: bool, | |
810 | /// `cdc` (publishes click-throughs as CDC events). | |
811 | pub sink: String, | |
812 | } | |
813 | ||
814 | impl Default for SearchUiAnalyticsConfig { | |
815 | 48 | fn default() -> Self { |
816 | 48 | Self { |
817 | 48 | enabled: false, |
818 | 48 | sink: "cdc".into(), |
819 | 48 | } |
820 | 48 | } |
821 | } |
Line | Count | Source |
1 | //! Layered configuration loading: file → env-var overrides → CLI overrides. | |
2 | ||
3 | use super::{ConfigError, MiroirConfig}; | |
4 | ||
5 | // The local `config` module shadows the external `config` crate. | |
6 | // Use a crate-qualified path to reach the external config crate. | |
7 | use ::config as ext_config; | |
8 | use serde_yaml; | |
9 | ||
10 | /// Default config file paths to search (in order). | |
11 | const CONFIG_SEARCH_PATHS: &[&str] = &[ | |
12 | "miroir.yaml", | |
13 | "miroir.yml", | |
14 | "config/miroir.yaml", | |
15 | "/etc/miroir/config.yaml", | |
16 | ]; | |
17 | ||
18 | /// Environment variable prefix for overrides. | |
19 | const ENV_PREFIX: &str = "MIROIR"; | |
20 | ||
21 | /// Load configuration using layered approach: | |
22 | /// 1. Search for config file in default paths | |
23 | /// 2. Apply environment variable overrides (`MIROIR_*`) | |
24 | /// 3. Returns validated config | |
25 | 0 | pub fn load() -> Result<MiroirConfig, ConfigError> { |
26 | 0 | let mut builder = ext_config::Config::builder(); |
27 | ||
28 | 0 | builder = builder.add_source(ext_config::Config::try_from(&MiroirConfig::default())?); |
29 | ||
30 | 0 | for path in CONFIG_SEARCH_PATHS { |
31 | 0 | if std::path::Path::new(path).exists() { |
32 | 0 | builder = builder.add_source(ext_config::File::with_name(path)); |
33 | 0 | break; |
34 | 0 | } |
35 | } | |
36 | ||
37 | 0 | builder = builder.add_source( |
38 | 0 | ext_config::Environment::with_prefix(ENV_PREFIX) |
39 | 0 | .separator("_") |
40 | 0 | .try_parsing(true), |
41 | ); | |
42 | ||
43 | 0 | let cfg: MiroirConfig = builder.build()?.try_deserialize()?; |
44 | 0 | cfg.validate()?; |
45 | 0 | Ok(cfg) |
46 | 0 | } |
47 | ||
48 | /// Load from a specific file path with env-var overrides applied. | |
49 | 0 | pub fn load_from(path: &std::path::Path) -> Result<MiroirConfig, ConfigError> { |
50 | 0 | let mut builder = ext_config::Config::builder(); |
51 | ||
52 | 0 | builder = builder.add_source(ext_config::Config::try_from(&MiroirConfig::default())?); |
53 | 0 | builder = builder.add_source(ext_config::File::with_name(path.to_string_lossy().as_ref())); |
54 | ||
55 | 0 | builder = builder.add_source( |
56 | 0 | ext_config::Environment::with_prefix(ENV_PREFIX) |
57 | 0 | .separator("_") |
58 | 0 | .try_parsing(true), |
59 | ); | |
60 | ||
61 | 0 | let cfg: MiroirConfig = builder.build()?.try_deserialize()?; |
62 | 0 | cfg.validate()?; |
63 | 0 | Ok(cfg) |
64 | 0 | } |
65 | ||
66 | /// Load from a YAML string (useful for testing). | |
67 | 12 | pub fn from_yaml(yaml: &str) -> Result<MiroirConfig, ConfigError> { |
68 | 12 | let |
69 | 10 | cfg.validate() |
70 | 6 | Ok(cfg) |
71 | 12 | } |
72 | ||
73 | #[cfg(test)] | |
74 | mod tests { | |
75 | use super::*; | |
76 | ||
77 | #[test] | |
78 | 2 | fn test_from_yaml_valid_config() { |
79 | 2 | let yaml = r#" |
80 | 2 | shards: 32 |
81 | 2 | replication_factor: 1 |
82 | 2 | cdc: |
83 | 2 | enabled: false |
84 | 2 | search_ui: |
85 | 2 | rate_limit: |
86 | 2 | backend: local |
87 | 2 | nodes: [] |
88 | 2 | "#; |
89 | 2 | let cfg = from_yaml(yaml).expect("should parse valid config"); |
90 | 2 | assert_eq!(cfg.shards, 32); |
91 | 2 | assert_eq!(cfg.replication_factor, 1); |
92 | 2 | } |
93 | ||
94 | #[test] | |
95 | 2 | fn test_from_yaml_with_nodes() { |
96 | 2 | let yaml = r#" |
97 | 2 | shards: 64 |
98 | 2 | replication_factor: 1 |
99 | 2 | replica_groups: 2 |
100 | 2 | task_store: |
101 | 2 | backend: redis |
102 | 2 | url: "redis://localhost:6379" |
103 | 2 | nodes: |
104 | 2 | - id: "node1" |
105 | 2 | address: "http://node1:7700" |
106 | 2 | replica_group: 0 |
107 | 2 | - id: "node2" |
108 | 2 | address: "http://node2:7700" |
109 | 2 | replica_group: 1 |
110 | 2 | "#; |
111 | 2 | let cfg = from_yaml(yaml).expect("should parse config with nodes"); |
112 | 2 | assert_eq!(cfg.nodes.len(), 2); |
113 | 2 | assert_eq!(cfg.nodes[0].id, "node1"); |
114 | 2 | assert_eq!(cfg.nodes[1].replica_group, 1); |
115 | 2 | } |
116 | ||
117 | #[test] | |
118 | 2 | fn test_from_yaml_invalid_yaml_fails() { |
119 | 2 | let yaml = r#" |
120 | 2 | shards: 32 |
121 | 2 | replication_factor: invalid |
122 | 2 | nodes: [] |
123 | 2 | "#; |
124 | 2 | let result = from_yaml(yaml); |
125 | 2 | assert!(result.is_err(), |
126 | 2 | } |
127 | ||
128 | #[test] | |
129 | 2 | fn test_from_yaml_validation_fails_on_ha_with_sqlite() { |
130 | 2 | let yaml = r#" |
131 | 2 | shards: 64 |
132 | 2 | replication_factor: 2 |
133 | 2 | nodes: [] |
134 | 2 | "#; |
135 | 2 | let result = from_yaml(yaml); |
136 | 2 | assert!(result.is_err(), |
137 | 2 | } |
138 | ||
139 | #[test] | |
140 | 2 | fn test_from_yaml_validation_fails_on_zero_shards() { |
141 | 2 | let yaml = r#" |
142 | 2 | shards: 0 |
143 | 2 | replication_factor: 1 |
144 | 2 | nodes: [] |
145 | 2 | "#; |
146 | 2 | let result = from_yaml(yaml); |
147 | 2 | assert!(result.is_err(), |
148 | 2 | } |
149 | ||
150 | #[test] | |
151 | 2 | fn test_from_yaml_with_all_sections() { |
152 | 2 | let yaml = r#" |
153 | 2 | shards: 64 |
154 | 2 | replication_factor: 1 |
155 | 2 | replica_groups: 2 |
156 | 2 | master_key: "test-key" |
157 | 2 | node_master_key: "node-key" |
158 | 2 | nodes: |
159 | 2 | - id: "node1" |
160 | 2 | address: "http://node1:7700" |
161 | 2 | replica_group: 0 |
162 | 2 | task_store: |
163 | 2 | backend: redis |
164 | 2 | url: "redis://localhost:6379" |
165 | 2 | admin: |
166 | 2 | enabled: true |
167 | 2 | api_key: "admin-key" |
168 | 2 | health: |
169 | 2 | interval_ms: 5000 |
170 | 2 | timeout_ms: 2000 |
171 | 2 | unhealthy_threshold: 3 |
172 | 2 | recovery_threshold: 2 |
173 | 2 | scatter: |
174 | 2 | node_timeout_ms: 5000 |
175 | 2 | retry_on_timeout: true |
176 | 2 | unavailable_shard_policy: partial |
177 | 2 | rebalancer: |
178 | 2 | auto_rebalance_on_recovery: true |
179 | 2 | max_concurrent_migrations: 4 |
180 | 2 | migration_timeout_s: 3600 |
181 | 2 | server: |
182 | 2 | port: 7700 |
183 | 2 | bind: "0.0.0.0" |
184 | 2 | max_body_bytes: 104857600 |
185 | 2 | leader_election: |
186 | 2 | enabled: true |
187 | 2 | "#; |
188 | 2 | let cfg = from_yaml(yaml).expect("should parse full config"); |
189 | 2 | assert_eq!(cfg.shards, 64); |
190 | 2 | assert_eq!(cfg.master_key, "test-key"); |
191 | 2 | assert_eq!(cfg.admin.api_key, "admin-key"); |
192 | 2 | assert_eq!(cfg.health.interval_ms, 5000); |
193 | 2 | assert_eq!(cfg.scatter.node_timeout_ms, 5000); |
194 | 2 | } |
195 | } |
Line | Count | Source |
1 | use crate::config::{ConfigError, MiroirConfig}; | |
2 | ||
3 | 22 | pub fn validate(cfg: &MiroirConfig) -> Result<(), ConfigError> { |
4 | // replication_factor > 1 requires redis backend for HA | |
5 | 22 | if cfg.replication_factor > 1 && |
6 | 4 | return Err(ConfigError::Validation( |
7 | 4 | "replication_factor > 1 requires task_store.backend = 'redis' (SQLite is single-writer)".into(), |
8 | 4 | )); |
9 | 18 | } |
10 | ||
11 | // replica_groups > 1 requires redis backend | |
12 | 18 | if cfg.replica_groups > 1 && |
13 | 0 | return Err(ConfigError::Validation( |
14 | 0 | "replica_groups > 1 requires task_store.backend = 'redis' (SQLite is single-writer)" |
15 | 0 | .into(), |
16 | 0 | )); |
17 | 18 | } |
18 | ||
19 | // Nodes must belong to a valid replica group | |
20 | 18 | if cfg.replica_groups > 0 { |
21 | 32 | for |
22 | 16 | if node.replica_group >= cfg.replica_groups { |
23 | 2 | return Err(ConfigError::Validation(format!( |
24 | 2 | "node '{}' has replica_group={} but only {} groups exist (0..{})", |
25 | 2 | node.id, |
26 | 2 | node.replica_group, |
27 | 2 | cfg.replica_groups, |
28 | 2 | cfg.replica_groups - 1 |
29 | 2 | ))); |
30 | 14 | } |
31 | } | |
32 | 0 | } |
33 | ||
34 | // Node IDs must be unique | |
35 | 16 | let mut seen_ids = std::collections::HashSet::new(); |
36 | 28 | for |
37 | 14 | if !seen_ids.insert(&node.id) { |
38 | 2 | return Err(ConfigError::Validation(format!( |
39 | 2 | "duplicate node id: '{}'", |
40 | 2 | node.id |
41 | 2 | ))); |
42 | 12 | } |
43 | } | |
44 | ||
45 | // HPA enabled requires redis backend | |
46 | 14 | if cfg.hpa.enabled && |
47 | 0 | return Err(ConfigError::Validation( |
48 | 0 | "hpa.enabled = true requires task_store.backend = 'redis'".into(), |
49 | 0 | )); |
50 | 14 | } |
51 | ||
52 | // Search UI scoped_key timing validation | |
53 | 14 | if cfg.search_ui.enabled { |
54 | 14 | let max_age = cfg.search_ui.scoped_key_max_age_days; |
55 | 14 | let rotate_before = cfg.search_ui.scoped_key_rotate_before_expiry_days; |
56 | 14 | if rotate_before >= max_age { |
57 | 2 | return Err(ConfigError::Validation(format!( |
58 | 2 | "search_ui.scoped_key_rotate_before_expiry_days ({rotate_before}) must be strictly less than scoped_key_max_age_days ({max_age})" |
59 | 2 | ))); |
60 | 12 | } |
61 | 0 | } |
62 | ||
63 | // CDC overflow = redis requires redis backend | |
64 | 12 | if cfg.cdc.enabled && |
65 | 2 | return Err(ConfigError::Validation( |
66 | 2 | "cdc.buffer.overflow = 'redis' requires task_store.backend = 'redis'".into(), |
67 | 2 | )); |
68 | 10 | } |
69 | ||
70 | // Search UI rate_limit.backend = redis requires redis task store (when multi-pod) | |
71 | 10 | if cfg.search_ui.enabled |
72 | 10 | && cfg.search_ui.rate_limit.backend == "redis" |
73 | 6 | && cfg.task_store.backend != "redis" |
74 | { | |
75 | 0 | return Err(ConfigError::Validation( |
76 | 0 | "search_ui.rate_limit.backend = 'redis' requires task_store.backend = 'redis'".into(), |
77 | 0 | )); |
78 | 10 | } |
79 | ||
80 | // Leader election should be enabled when replica_groups > 1 | |
81 | 10 | if cfg.replica_groups > 1 && |
82 | 0 | return Err(ConfigError::Validation( |
83 | 0 | "leader_election.enabled must be true when replica_groups > 1".into(), |
84 | 0 | )); |
85 | 10 | } |
86 | ||
87 | // Tenant affinity dedicated_groups must be within valid range | |
88 | 10 | if cfg.tenant_affinity.enabled { |
89 | 10 | for |
90 | 0 | if *g >= cfg.replica_groups { |
91 | 0 | return Err(ConfigError::Validation(format!( |
92 | 0 | "tenant_affinity.dedicated_groups contains {} but only {} groups (0..{})", |
93 | 0 | g, |
94 | 0 | cfg.replica_groups, |
95 | 0 | cfg.replica_groups - 1 |
96 | 0 | ))); |
97 | 0 | } |
98 | } | |
99 | 10 | for ( |
100 | 0 | if *group >= cfg.replica_groups { |
101 | 0 | return Err(ConfigError::Validation(format!( |
102 | 0 | "tenant_affinity.static_map: tenant '{}' maps to group {} but only {} groups (0..{})", |
103 | 0 | tenant, |
104 | 0 | group, |
105 | 0 | cfg.replica_groups, |
106 | 0 | cfg.replica_groups - 1 |
107 | 0 | ))); |
108 | 0 | } |
109 | } | |
110 | 0 | } |
111 | ||
112 | // Shadow targets must have valid sample_rate | |
113 | 10 | if cfg.shadow.enabled { |
114 | 10 | for |
115 | 0 | if target.sample_rate <= 0.0 || target.sample_rate > 1.0 { |
116 | 0 | return Err(ConfigError::Validation(format!( |
117 | 0 | "shadow target '{}' has invalid sample_rate={} (must be 0 < rate <= 1)", |
118 | 0 | target.name, target.sample_rate |
119 | 0 | ))); |
120 | 0 | } |
121 | } | |
122 | 0 | } |
123 | ||
124 | // Server port must be non-zero | |
125 | 10 | if cfg.server.port == 0 { |
126 | 0 | return Err(ConfigError::Validation( |
127 | 0 | "server.port must be non-zero".into(), |
128 | 0 | )); |
129 | 10 | } |
130 | ||
131 | // shards must be non-zero | |
132 | 10 | if cfg.shards == 0 { |
133 | 2 | return Err(ConfigError::Validation("shards must be non-zero".into())); |
134 | 8 | } |
135 | ||
136 | // replication_factor must be > 0 | |
137 | 8 | if cfg.replication_factor == 0 { |
138 | 0 | return Err(ConfigError::Validation( |
139 | 0 | "replication_factor must be > 0".into(), |
140 | 0 | )); |
141 | 8 | } |
142 | ||
143 | 8 | Ok(()) |
144 | 22 | } |
Line | Count | Source |
1 | //! Result merger: combines shard results into a single response. | |
2 | ||
3 | use crate::Result; | |
4 | use serde::{Deserialize, Serialize}; | |
5 | use serde_json::Value; | |
6 | use std::collections::{BTreeMap, BinaryHeap}; | |
7 | ||
8 | /// Input to the merge function. | |
9 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
10 | pub struct MergeInput { | |
11 | /// One page of hits per node in the covering set. | |
12 | pub shard_hits: Vec<ShardHitPage>, | |
13 | ||
14 | /// Offset to apply after merge. | |
15 | pub offset: usize, | |
16 | ||
17 | /// Limit to apply after merge. | |
18 | pub limit: usize, | |
19 | ||
20 | /// Whether the client requested the ranking score in the response. | |
21 | pub client_requested_score: bool, | |
22 | ||
23 | /// Facet names to include (if None, include all). | |
24 | pub facets: Option<Vec<String>>, | |
25 | } | |
26 | ||
27 | /// A page of hits from a single shard. | |
28 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
29 | pub struct ShardHitPage { | |
30 | /// Raw JSON response from the node. | |
31 | pub body: Value, | |
32 | ||
33 | /// Whether this shard succeeded. | |
34 | pub success: bool, | |
35 | } | |
36 | ||
37 | /// Merged search result. | |
38 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
39 | pub struct MergedSearchResult { | |
40 | /// Merged hits (globally sorted, offset/limit applied). | |
41 | pub hits: Vec<Value>, | |
42 | ||
43 | /// Aggregated facets. | |
44 | pub facets: Value, | |
45 | ||
46 | /// Estimated total hits (sum of shard totals). | |
47 | pub estimated_total_hits: u64, | |
48 | ||
49 | /// Processing time in milliseconds. | |
50 | pub processing_time_ms: u64, | |
51 | ||
52 | /// Whether the response is degraded (some shards failed). | |
53 | pub degraded: bool, | |
54 | } | |
55 | ||
56 | /// Merge search results from multiple shards. | |
57 | /// | |
58 | /// This is the primary entry point for result merging. It: | |
59 | /// 1. Collects all hits with scores from all shards | |
60 | /// 2. Sorts globally by `_rankingScore` descending | |
61 | /// 3. Applies offset + limit after merge | |
62 | /// 4. Strips `_rankingScore` if client didn't request it | |
63 | /// 5. Always strips `_miroir_*` reserved fields | |
64 | /// 6. Sums facet counts across shards | |
65 | /// 7. Sums `estimatedTotalHits` across shards | |
66 | /// 8. Takes max `processingTimeMs` across shards | |
67 | /// | |
68 | /// Uses a binary min-heap to avoid keeping all hits in RAM when fan-out is large. | |
69 | /// Uses BTreeMap for stable facet serialization (deterministic JSON output). | |
70 | 54 | pub fn merge(input: MergeInput) -> MergedSearchResult { |
71 | // Filter to only successful responses | |
72 | 54 | let successful_shards: Vec<_> = input |
73 | 54 | .shard_hits |
74 | 54 | .iter() |
75 | 54 | .filter(|s| s.success) |
76 | 54 | .collect(); |
77 | ||
78 | // Check if any shards failed (degraded mode) | |
79 | 54 | let degraded = successful_shards.len() < input.shard_hits.len(); |
80 | ||
81 | // Collect all hits with their ranking scores and primary keys | |
82 | 54 | let mut all_hits: Vec<HitWithScore> = Vec::new(); |
83 | ||
84 | 128 | for |
85 | 74 | if let Some(hits) = shard.body.get("hits").and_then(|h| h.as_array()) { |
86 | 924 | for |
87 | 850 | let score = hit |
88 | 850 | .get("_rankingScore") |
89 | 850 | .and_then(|s| s.as_f64()) |
90 | 850 | .unwrap_or(0.0); |
91 | ||
92 | 850 | let primary_key = hit |
93 | 850 | .get("id") |
94 | 850 | .and_then(|id| id.as_str()) |
95 | 850 | .unwrap_or("") |
96 | 850 | .to_string(); |
97 | ||
98 | 850 | all_hits.push(HitWithScore { |
99 | 850 | score, |
100 | 850 | primary_key, |
101 | 850 | hit: hit.clone(), |
102 | 850 | }); |
103 | } | |
104 | 0 | } |
105 | } | |
106 | ||
107 | // Use a min-heap of size offset + limit to avoid keeping all hits in RAM | |
108 | 54 | let heap_size = input.offset + input.limit; |
109 | 54 | let top_hits = if all_hits.len() > heap_size * 2 { |
110 | // Only use heap optimization if we have significantly more hits than needed | |
111 | // Use BinaryHeap with Reverse to get min-heap behavior | |
112 | // Reverse<T> makes the "smallest" T become the "largest" Reverse<T> | |
113 | // So in BinaryHeap (max-heap), the smallest T ends up at the top | |
114 | use std::cmp::Reverse; | |
115 | ||
116 | // First, define a wrapper with natural ascending order by score | |
117 | #[derive(Debug, Clone, PartialEq, Eq)] | |
118 | struct AscendingHit(HitWithScore); | |
119 | ||
120 | impl PartialOrd for AscendingHit { | |
121 | 968 | fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
122 | 968 | Some(self.cmp(other)) |
123 | 968 | } |
124 | } | |
125 | ||
126 | impl Ord for AscendingHit { | |
127 | 968 | fn cmp(&self, other: &Self) -> std::cmp::Ordering { |
128 | // Ascending order: lower score first, then primary key ascending | |
129 | 968 | match self.0.score.partial_cmp(&other.0.score) { |
130 | Some(std::cmp::Ordering::Equal) | None => { | |
131 | 0 | self.0.primary_key.cmp(&other.0.primary_key) |
132 | } | |
133 | 968 | Some(ord) => ord, |
134 | } | |
135 | 968 | } |
136 | } | |
137 | ||
138 | // Reverse<AscendingHit> in BinaryHeap gives us a min-heap by score | |
139 | // The lowest score (worst) will be at the top for easy replacement | |
140 | 6 | let mut heap: BinaryHeap<Reverse<AscendingHit>> = BinaryHeap::new(); |
141 | ||
142 | 406 | for |
143 | 400 | if heap.len() < heap_size { |
144 | 80 | heap.push(Reverse(AscendingHit(hit))); |
145 | 80 | } else { |
146 | // Peek at the smallest (worst) hit in our top-k | |
147 | 320 | if let Some(Reverse(worst)) = heap.peek() { |
148 | // If current hit is better than our worst, replace it | |
149 | // Compare using the score directly | |
150 | 320 | if hit.score > worst.0.score || |
151 | 180 | ( |
152 | 180 | heap.pop(); |
153 | 180 | heap.push(Reverse(AscendingHit(hit))); |
154 | 180 |
|
155 | 0 | } |
156 | } | |
157 | } | |
158 | ||
159 | // Convert to sorted vector (descending by score) | |
160 | 6 | let mut result: Vec<_> = heap.into_iter().map(|r| (r.0).0).collect(); |
161 | 500 |
|
162 | 6 | result |
163 | } else { | |
164 | // For smaller result sets, just sort directly | |
165 | 436 |
|
166 | 48 | all_hits |
167 | }; | |
168 | ||
169 | // Apply offset and limit | |
170 | 54 | let page: Vec<Value> = top_hits |
171 | 54 | .into_iter() |
172 | 54 | .skip(input.offset) |
173 | 54 | .take(input.limit) |
174 | 260 | . |
175 | 260 | let hit = &mut hit_with_score.hit; |
176 | ||
177 | // Strip all _miroir_* fields (always removed) | |
178 | 260 | if let Some(obj) = hit.as_object_mut() { |
179 | 784 |
|
180 | 0 | } |
181 | ||
182 | // Strip _rankingScore if client didn't request it | |
183 | 260 | if !input.client_requested_score { |
184 | 252 | if let Some(obj) = hit.as_object_mut() { |
185 | 252 | obj.remove("_rankingScore"); |
186 | 252 |
|
187 | 8 | } |
188 | ||
189 | 260 | hit.clone() |
190 | 260 | }) |
191 | 54 | .collect(); |
192 | ||
193 | // Aggregate facets across all shards | |
194 | 54 | let facets = merge_facets(&successful_shards, input.facets.as_deref()); |
195 | ||
196 | // Sum estimated total hits | |
197 | 54 | let estimated_total_hits: u64 = successful_shards |
198 | 54 | .iter() |
199 | 74 | . |
200 | 54 | .sum(); |
201 | ||
202 | // Max processing time across all shards | |
203 | 54 | let processing_time_ms: u64 = successful_shards |
204 | 54 | .iter() |
205 | 74 | . |
206 | 54 | .max() |
207 | 54 | .unwrap_or(0); |
208 | ||
209 | 54 | MergedSearchResult { |
210 | 54 | hits: page, |
211 | 54 | facets, |
212 | 54 | estimated_total_hits, |
213 | 54 | processing_time_ms, |
214 | 54 | degraded, |
215 | 54 | } |
216 | 54 | } |
217 | ||
218 | /// A hit with its ranking score for sorting. | |
219 | #[derive(Debug, Clone)] | |
220 | struct HitWithScore { | |
221 | score: f64, | |
222 | primary_key: String, | |
223 | hit: Value, | |
224 | } | |
225 | ||
226 | impl PartialEq for HitWithScore { | |
227 | 0 | fn eq(&self, other: &Self) -> bool { |
228 | 0 | self.score == other.score && self.primary_key == other.primary_key |
229 | 0 | } |
230 | } | |
231 | ||
232 | impl Eq for HitWithScore {} | |
233 | ||
234 | impl PartialOrd for HitWithScore { | |
235 | 0 | fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> { |
236 | 0 | Some(self.cmp(other)) |
237 | 0 | } |
238 | } | |
239 | ||
240 | impl HitWithScore { | |
241 | /// Compare for descending order (score desc, primary key asc). | |
242 | 936 | fn descending_cmp(&self, other: &Self) -> std::cmp::Ordering { |
243 | // Sort by score descending, then by primary key ascending for tie-breaking | |
244 | 936 | match other.score.partial_cmp(&self.score) { |
245 | Some(std::cmp::Ordering::Equal) | None => { | |
246 | // On equal scores (or NaN), fall back to primary key ascending | |
247 | 6 | self.primary_key.cmp(&other.primary_key) |
248 | } | |
249 | 930 | Some(ord) => ord, |
250 | } | |
251 | 936 | } |
252 | } | |
253 | ||
254 | impl Ord for HitWithScore { | |
255 | 0 | fn cmp(&self, other: &Self) -> std::cmp::Ordering { |
256 | // Sort by score descending, then by primary key ascending for tie-breaking | |
257 | 0 | match self.score.partial_cmp(&other.score) { |
258 | Some(std::cmp::Ordering::Equal) | None => { | |
259 | // On equal scores (or NaN), fall back to primary key | |
260 | // Reversed for descending order | |
261 | 0 | other.primary_key.cmp(&self.primary_key) |
262 | } | |
263 | 0 | Some(ord) => ord.reverse(), |
264 | } | |
265 | 0 | } |
266 | } | |
267 | ||
268 | /// Result merger: combines responses from multiple shards. | |
269 | pub trait Merger: Send + Sync { | |
270 | /// Merge search results from multiple shards. | |
271 | /// | |
272 | /// Takes the raw JSON responses from each shard and produces | |
273 | /// a merged result with global sorting, offset/limit applied, | |
274 | /// and facet aggregation. | |
275 | fn merge( | |
276 | &self, | |
277 | shard_responses: Vec<ShardResponse>, | |
278 | offset: usize, | |
279 | limit: usize, | |
280 | client_requested_score: bool, | |
281 | ) -> Result<MergedResult>; | |
282 | } | |
283 | ||
284 | /// Response from a single shard. | |
285 | #[derive(Debug, Clone)] | |
286 | pub struct ShardResponse { | |
287 | /// Shard identifier. | |
288 | pub shard_id: u32, | |
289 | ||
290 | /// Raw JSON response from the node. | |
291 | pub body: Value, | |
292 | ||
293 | /// Whether this shard succeeded. | |
294 | pub success: bool, | |
295 | } | |
296 | ||
297 | /// Merged search result (legacy, use MergedSearchResult instead). | |
298 | #[derive(Debug, Clone)] | |
299 | pub struct MergedResult { | |
300 | /// Merged hits (globally sorted, offset/limit applied). | |
301 | pub hits: Vec<Value>, | |
302 | ||
303 | /// Aggregated facets. | |
304 | pub facets: Value, | |
305 | ||
306 | /// Estimated total hits (sum of shard totals). | |
307 | pub total_hits: u64, | |
308 | ||
309 | /// Processing time in milliseconds. | |
310 | pub processing_time_ms: u64, | |
311 | ||
312 | /// Whether the response is degraded (some shards failed). | |
313 | pub degraded: bool, | |
314 | } | |
315 | ||
316 | /// Default implementation of Merger. | |
317 | #[derive(Debug, Clone, Default)] | |
318 | pub struct MergerImpl; | |
319 | ||
320 | impl Merger for MergerImpl { | |
321 | 48 | fn merge( |
322 | 48 | &self, |
323 | 48 | shard_responses: Vec<ShardResponse>, |
324 | 48 | offset: usize, |
325 | 48 | limit: usize, |
326 | 48 | client_requested_score: bool, |
327 | 48 | ) -> Result<MergedResult> { |
328 | // Convert ShardResponse to ShardHitPage | |
329 | 48 | let shard_hits: Vec<ShardHitPage> = shard_responses |
330 | 48 | .into_iter() |
331 | 48 | .map(|sr| ShardHitPage { |
332 | 66 | body: sr.body, |
333 | 66 | success: sr.success, |
334 | 66 | }) |
335 | 48 | .collect(); |
336 | ||
337 | // Call the new merge function | |
338 | 48 | let input = MergeInput { |
339 | 48 | shard_hits, |
340 | 48 | offset, |
341 | 48 | limit, |
342 | 48 | client_requested_score, |
343 | 48 | facets: None, |
344 | 48 | }; |
345 | ||
346 | 48 | let result = merge(input); |
347 | ||
348 | 48 | Ok(MergedResult { |
349 | 48 | hits: result.hits, |
350 | 48 | facets: result.facets, |
351 | 48 | total_hits: result.estimated_total_hits, |
352 | 48 | processing_time_ms: result.processing_time_ms, |
353 | 48 | degraded: result.degraded, |
354 | 48 | }) |
355 | 48 | } |
356 | } | |
357 | ||
358 | /// Merge facet distributions from all shards. | |
359 | /// | |
360 | /// Facets are nested objects like `{"color": {"red": 10, "blue": 5}}`. | |
361 | /// We sum counts for each facet value across all shards. | |
362 | /// | |
363 | /// Uses BTreeMap for stable, deterministic serialization. | |
364 | 54 | fn merge_facets(shards: &[&ShardHitPage], facet_filter: Option<&[String]>) -> Value { |
365 | 54 | let mut merged_facets: BTreeMap<String, BTreeMap<String, u64>> = BTreeMap::new(); |
366 | ||
367 | 128 | for |
368 | 74 | if let Some( |
369 | 74 | .body |
370 | 74 | .get("facetDistribution") |
371 | 74 | .and_then(|f| |
372 | { | |
373 | 100 | for ( |
374 | // Apply facet filter if provided | |
375 | 34 | if let Some( |
376 | 4 | if !filter.contains(&facet_name) { |
377 | 2 | continue; |
378 | 2 | } |
379 | 30 | } |
380 | ||
381 | 32 | if let Some(values_obj) = facet_values.as_object() { |
382 | 32 | let entry = merged_facets.entry(facet_name.clone()).or_default(); |
383 | ||
384 | 88 | for ( |
385 | 56 | let count_val = count.as_u64().unwrap_or(0); |
386 | 56 | *entry.entry(value.clone()).or_insert(0) += count_val; |
387 | 56 | } |
388 | 0 | } |
389 | } | |
390 | 8 | } |
391 | } | |
392 | ||
393 | // Convert back to JSON structure (BTreeMap ensures stable key order) | |
394 | 54 | let result: serde_json::Map<String, Value> = merged_facets |
395 | 54 | .into_iter() |
396 | 54 | .map(|(facet_name, values)| |
397 | 18 | let values_obj: serde_json::Map<String, Value> = values |
398 | 18 | .into_iter() |
399 | 44 | . |
400 | 18 | .collect(); |
401 | 18 | (facet_name, Value::Object(values_obj)) |
402 | 18 | }) |
403 | 54 | .collect(); |
404 | ||
405 | 54 | Value::Object(result) |
406 | 54 | } |
407 | ||
408 | /// Stub implementation that returns empty results. | |
409 | #[derive(Debug, Clone, Default)] | |
410 | pub struct StubMerger; | |
411 | ||
412 | impl Merger for StubMerger { | |
413 | 0 | fn merge( |
414 | 0 | &self, |
415 | 0 | _shard_responses: Vec<ShardResponse>, |
416 | 0 | _offset: usize, |
417 | 0 | _limit: usize, |
418 | 0 | _client_requested_score: bool, |
419 | 0 | ) -> Result<MergedResult> { |
420 | 0 | Ok(MergedResult { |
421 | 0 | hits: Vec::new(), |
422 | 0 | facets: serde_json::json!({}), |
423 | 0 | total_hits: 0, |
424 | 0 | processing_time_ms: 0, |
425 | 0 | degraded: false, |
426 | 0 | }) |
427 | 0 | } |
428 | } | |
429 | ||
430 | #[cfg(test)] | |
431 | mod tests { | |
432 | use super::*; | |
433 | ||
434 | 348 | fn create_hit(score: f64, id: &str) -> Value { |
435 | 348 | serde_json::json!({ |
436 | 348 | "id": id, |
437 | 348 | "_rankingScore": score, |
438 | 348 | "_miroir_shard": 0, |
439 | }) | |
440 | 348 | } |
441 | ||
442 | 46 | fn create_shard_response(shard_id: u32, hits: Vec<Value>, total: u64) -> ShardResponse { |
443 | 46 | ShardResponse { |
444 | 46 | shard_id, |
445 | 46 | body: serde_json::json!({ |
446 | 46 | "hits": hits, |
447 | 46 | "estimatedTotalHits": total, |
448 | 46 | "processingTimeMs": 10, |
449 | 46 | "facetDistribution": {}, |
450 | 46 | }), |
451 | 46 | success: true, |
452 | 46 | } |
453 | 46 | } |
454 | ||
455 | #[test] | |
456 | 2 | fn test_global_sort_by_ranking_score() { |
457 | 2 | let merger = MergerImpl; |
458 | ||
459 | 2 | let hits1 = vec![ |
460 | 2 | create_hit(0.5, "doc1"), |
461 | 2 | create_hit(0.9, "doc2"), |
462 | 2 | create_hit(0.3, "doc3"), |
463 | ]; | |
464 | 2 | let hits2 = vec![ |
465 | 2 | create_hit(0.7, "doc4"), |
466 | 2 | create_hit(0.1, "doc5"), |
467 | 2 | create_hit(0.8, "doc6"), |
468 | ]; | |
469 | ||
470 | 2 | let shards = vec![ |
471 | 2 | create_shard_response(0, hits1, 3), |
472 | 2 | create_shard_response(1, hits2, 3), |
473 | ]; | |
474 | ||
475 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
476 | ||
477 | // Check global ordering: should be doc2(0.9), doc6(0.8), doc4(0.7), doc1(0.5), doc3(0.3), doc5(0.1) | |
478 | 2 | assert_eq!(result.hits.len(), 6); |
479 | 2 | assert_eq!(result.hits[0]["id"], "doc2"); |
480 | 2 | assert_eq!(result.hits[1]["id"], "doc6"); |
481 | 2 | assert_eq!(result.hits[2]["id"], "doc4"); |
482 | 2 | assert_eq!(result.hits[3]["id"], "doc1"); |
483 | 2 | assert_eq!(result.hits[4]["id"], "doc3"); |
484 | 2 | assert_eq!(result.hits[5]["id"], "doc5"); |
485 | 2 | } |
486 | ||
487 | #[test] | |
488 | 2 | fn test_offset_and_limit_applied_after_merge() { |
489 | 2 | let merger = MergerImpl; |
490 | ||
491 | 2 | let hits = vec![ |
492 | 2 | create_hit(0.9, "doc1"), |
493 | 2 | create_hit(0.8, "doc2"), |
494 | 2 | create_hit(0.7, "doc3"), |
495 | 2 | create_hit(0.6, "doc4"), |
496 | 2 | create_hit(0.5, "doc5"), |
497 | ]; | |
498 | ||
499 | 2 | let shards = vec![create_shard_response(0, hits, 5)]; |
500 | ||
501 | 2 | let result = merger.merge(shards, 2, 2, false).unwrap(); |
502 | ||
503 | // Should skip first 2, take next 2 | |
504 | 2 | assert_eq!(result.hits.len(), 2); |
505 | 2 | assert_eq!(result.hits[0]["id"], "doc3"); |
506 | 2 | assert_eq!(result.hits[1]["id"], "doc4"); |
507 | 2 | } |
508 | ||
509 | #[test] | |
510 | 2 | fn test_ranking_score_stripped_when_not_requested() { |
511 | 2 | let merger = MergerImpl; |
512 | ||
513 | 2 | let hits = vec![create_hit(0.9, "doc1"), create_hit(0.8, "doc2")]; |
514 | ||
515 | 2 | let shards = vec![create_shard_response(0, hits, 2)]; |
516 | ||
517 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
518 | ||
519 | // _rankingScore should be stripped | |
520 | 2 | assert!(result.hits[0].get("_rankingScore").is_none()); |
521 | 2 | assert!(result.hits[1].get("_rankingScore").is_none()); |
522 | 2 | } |
523 | ||
524 | #[test] | |
525 | 2 | fn test_ranking_score_included_when_requested() { |
526 | 2 | let merger = MergerImpl; |
527 | ||
528 | 2 | let hits = vec![create_hit(0.9, "doc1"), create_hit(0.8, "doc2")]; |
529 | ||
530 | 2 | let shards = vec![create_shard_response(0, hits, 2)]; |
531 | ||
532 | 2 | let result = merger.merge(shards, 0, 10, true).unwrap(); |
533 | ||
534 | // _rankingScore should be present | |
535 | 2 | assert_eq!(result.hits[0]["_rankingScore"], 0.9); |
536 | 2 | assert_eq!(result.hits[1]["_rankingScore"], 0.8); |
537 | 2 | } |
538 | ||
539 | #[test] | |
540 | 2 | fn test_miroir_shard_always_stripped() { |
541 | 2 | let merger = MergerImpl; |
542 | ||
543 | 2 | let hits = vec![create_hit(0.9, "doc1")]; |
544 | ||
545 | 2 | let shards = vec![create_shard_response(0, hits, 1)]; |
546 | ||
547 | 2 | let result = merger.merge(shards, 0, 10, true).unwrap(); |
548 | ||
549 | // _miroir_shard should always be stripped | |
550 | 2 | assert!(result.hits[0].get("_miroir_shard").is_none()); |
551 | 2 | } |
552 | ||
553 | #[test] | |
554 | 2 | fn test_facet_counts_summed_across_shards() { |
555 | 2 | let merger = MergerImpl; |
556 | ||
557 | 2 | let shard1 = serde_json::json!({ |
558 | 2 | "hits": [], |
559 | 2 | "estimatedTotalHits": 0, |
560 | 2 | "processingTimeMs": 10, |
561 | 2 | "facetDistribution": { |
562 | 2 | "color": { |
563 | 2 | "red": 10, |
564 | 2 | "blue": 5, |
565 | }, | |
566 | 2 | "size": { |
567 | 2 | "large": 8, |
568 | } | |
569 | } | |
570 | }); | |
571 | ||
572 | 2 | let shard2 = serde_json::json!({ |
573 | 2 | "hits": [], |
574 | 2 | "estimatedTotalHits": 0, |
575 | 2 | "processingTimeMs": 15, |
576 | 2 | "facetDistribution": { |
577 | 2 | "color": { |
578 | 2 | "red": 7, |
579 | 2 | "green": 3, |
580 | }, | |
581 | 2 | "size": { |
582 | 2 | "large": 4, |
583 | 2 | "small": 6, |
584 | } | |
585 | } | |
586 | }); | |
587 | ||
588 | 2 | let shards = vec![ |
589 | 2 | ShardResponse { |
590 | 2 | shard_id: 0, |
591 | 2 | body: shard1, |
592 | 2 | success: true, |
593 | 2 | }, |
594 | 2 | ShardResponse { |
595 | 2 | shard_id: 1, |
596 | 2 | body: shard2, |
597 | 2 | success: true, |
598 | 2 | }, |
599 | ]; | |
600 | ||
601 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
602 | ||
603 | 2 | let facets = result.facets; |
604 | 2 | let color = facets.get("color").unwrap().as_object().unwrap(); |
605 | 2 | assert_eq!(color.get("red").unwrap().as_u64().unwrap(), 17); // 10 + 7 |
606 | 2 | assert_eq!(color.get("blue").unwrap().as_u64().unwrap(), 5); // only in shard1 |
607 | 2 | assert_eq!(color.get("green").unwrap().as_u64().unwrap(), 3); // only in shard2 |
608 | ||
609 | 2 | let size = facets.get("size").unwrap().as_object().unwrap(); |
610 | 2 | assert_eq!(size.get("large").unwrap().as_u64().unwrap(), 12); // 8 + 4 |
611 | 2 | assert_eq!(size.get("small").unwrap().as_u64().unwrap(), 6); // only in shard2 |
612 | 2 | } |
613 | ||
614 | #[test] | |
615 | 2 | fn test_estimated_total_hits_summed() { |
616 | 2 | let merger = MergerImpl; |
617 | ||
618 | 2 | let shards = vec![ |
619 | 2 | create_shard_response(0, vec![], 100), |
620 | 2 | create_shard_response(1, vec![], 150), |
621 | 2 | create_shard_response(2, vec![], 75), |
622 | ]; | |
623 | ||
624 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
625 | ||
626 | 2 | assert_eq!(result.total_hits, 325); // 100 + 150 + 75 |
627 | 2 | } |
628 | ||
629 | #[test] | |
630 | 2 | fn test_processing_time_max_across_shards() { |
631 | 2 | let merger = MergerImpl; |
632 | ||
633 | 2 | let shard1 = serde_json::json!({ |
634 | 2 | "hits": [], |
635 | 2 | "estimatedTotalHits": 0, |
636 | 2 | "processingTimeMs": 10, |
637 | }); | |
638 | ||
639 | 2 | let shard2 = serde_json::json!({ |
640 | 2 | "hits": [], |
641 | 2 | "estimatedTotalHits": 0, |
642 | 2 | "processingTimeMs": 25, |
643 | }); | |
644 | ||
645 | 2 | let shard3 = serde_json::json!({ |
646 | 2 | "hits": [], |
647 | 2 | "estimatedTotalHits": 0, |
648 | 2 | "processingTimeMs": 15, |
649 | }); | |
650 | ||
651 | 2 | let shards = vec![ |
652 | 2 | ShardResponse { |
653 | 2 | shard_id: 0, |
654 | 2 | body: shard1, |
655 | 2 | success: true, |
656 | 2 | }, |
657 | 2 | ShardResponse { |
658 | 2 | shard_id: 1, |
659 | 2 | body: shard2, |
660 | 2 | success: true, |
661 | 2 | }, |
662 | 2 | ShardResponse { |
663 | 2 | shard_id: 2, |
664 | 2 | body: shard3, |
665 | 2 | success: true, |
666 | 2 | }, |
667 | ]; | |
668 | ||
669 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
670 | ||
671 | 2 | assert_eq!(result.processing_time_ms, 25); // max(10, 25, 15) |
672 | 2 | } |
673 | ||
674 | #[test] | |
675 | 2 | fn test_degraded_flag_when_shard_fails() { |
676 | 2 | let merger = MergerImpl; |
677 | ||
678 | 2 | let hits = vec![create_hit(0.9, "doc1")]; |
679 | ||
680 | 2 | let shards = vec![ |
681 | 2 | create_shard_response(0, hits.clone(), 1), |
682 | 2 | ShardResponse { |
683 | 2 | shard_id: 1, |
684 | 2 | body: serde_json::json!({}), |
685 | 2 | success: false, |
686 | 2 | }, |
687 | ]; | |
688 | ||
689 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
690 | ||
691 | 2 | assert!(result.degraded); |
692 | 2 | assert_eq!(result.hits.len(), 1); |
693 | 2 | } |
694 | ||
695 | #[test] | |
696 | 2 | fn test_not_degraded_when_all_succeed() { |
697 | 2 | let merger = MergerImpl; |
698 | ||
699 | 2 | let shards = vec![ |
700 | 2 | create_shard_response(0, vec![], 0), |
701 | 2 | create_shard_response(1, vec![], 0), |
702 | ]; | |
703 | ||
704 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
705 | ||
706 | 2 | assert!(!result.degraded); |
707 | 2 | } |
708 | ||
709 | #[test] | |
710 | 2 | fn test_empty_shards_returns_empty_result() { |
711 | 2 | let merger = MergerImpl; |
712 | ||
713 | 2 | let result = merger.merge(vec![], 0, 10, false).unwrap(); |
714 | ||
715 | 2 | assert!(result.hits.is_empty()); |
716 | 2 | assert_eq!(result.total_hits, 0); |
717 | 2 | assert_eq!(result.facets.as_object().unwrap().len(), 0); |
718 | 2 | assert!(!result.degraded); |
719 | 2 | } |
720 | ||
721 | #[test] | |
722 | 2 | fn test_facet_keys_unique_to_one_shard_preserved() { |
723 | 2 | let merger = MergerImpl; |
724 | ||
725 | 2 | let shard1 = serde_json::json!({ |
726 | 2 | "hits": [], |
727 | 2 | "estimatedTotalHits": 0, |
728 | 2 | "processingTimeMs": 10, |
729 | 2 | "facetDistribution": { |
730 | 2 | "category": { |
731 | 2 | "electronics": 20, |
732 | 2 | "books": 15, |
733 | } | |
734 | } | |
735 | }); | |
736 | ||
737 | 2 | let shard2 = serde_json::json!({ |
738 | 2 | "hits": [], |
739 | 2 | "estimatedTotalHits": 0, |
740 | 2 | "processingTimeMs": 10, |
741 | 2 | "facetDistribution": { |
742 | 2 | "category": { |
743 | 2 | "clothing": 30, |
744 | 2 | "food": 10, |
745 | } | |
746 | } | |
747 | }); | |
748 | ||
749 | 2 | let shards = vec![ |
750 | 2 | ShardResponse { |
751 | 2 | shard_id: 0, |
752 | 2 | body: shard1, |
753 | 2 | success: true, |
754 | 2 | }, |
755 | 2 | ShardResponse { |
756 | 2 | shard_id: 1, |
757 | 2 | body: shard2, |
758 | 2 | success: true, |
759 | 2 | }, |
760 | ]; | |
761 | ||
762 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
763 | ||
764 | 2 | let category = result.facets.get("category").unwrap().as_object().unwrap(); |
765 | 2 | assert_eq!(category.get("electronics").unwrap().as_u64().unwrap(), 20); |
766 | 2 | assert_eq!(category.get("books").unwrap().as_u64().unwrap(), 15); |
767 | 2 | assert_eq!(category.get("clothing").unwrap().as_u64().unwrap(), 30); |
768 | 2 | assert_eq!(category.get("food").unwrap().as_u64().unwrap(), 10); |
769 | 2 | } |
770 | ||
771 | #[test] | |
772 | 2 | fn test_missing_facet_distribution_handled_gracefully() { |
773 | 2 | let merger = MergerImpl; |
774 | ||
775 | 2 | let shard_with_facets = serde_json::json!({ |
776 | 2 | "hits": [], |
777 | 2 | "estimatedTotalHits": 0, |
778 | 2 | "processingTimeMs": 10, |
779 | 2 | "facetDistribution": { |
780 | 2 | "color": {"red": 5} |
781 | } | |
782 | }); | |
783 | ||
784 | 2 | let shard_without_facets = serde_json::json!({ |
785 | 2 | "hits": [], |
786 | 2 | "estimatedTotalHits": 0, |
787 | 2 | "processingTimeMs": 10, |
788 | }); | |
789 | ||
790 | 2 | let shards = vec![ |
791 | 2 | ShardResponse { |
792 | 2 | shard_id: 0, |
793 | 2 | body: shard_with_facets, |
794 | 2 | success: true, |
795 | 2 | }, |
796 | 2 | ShardResponse { |
797 | 2 | shard_id: 1, |
798 | 2 | body: shard_without_facets, |
799 | 2 | success: true, |
800 | 2 | }, |
801 | ]; | |
802 | ||
803 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
804 | ||
805 | // Should still have facets from the shard that provided them | |
806 | 2 | let color = result.facets.get("color").unwrap().as_object().unwrap(); |
807 | 2 | assert_eq!(color.get("red").unwrap().as_u64().unwrap(), 5); |
808 | 2 | } |
809 | ||
810 | #[test] | |
811 | 2 | fn test_offset_exceeds_total_hits() { |
812 | 2 | let merger = MergerImpl; |
813 | ||
814 | 2 | let hits = vec![create_hit(0.9, "doc1"), create_hit(0.8, "doc2")]; |
815 | ||
816 | 2 | let shards = vec![create_shard_response(0, hits, 2)]; |
817 | ||
818 | 2 | let result = merger.merge(shards, 10, 10, false).unwrap(); |
819 | ||
820 | 2 | assert!(result.hits.is_empty()); |
821 | 2 | } |
822 | ||
823 | #[test] | |
824 | 2 | fn test_limit_exceeds_available_hits() { |
825 | 2 | let merger = MergerImpl; |
826 | ||
827 | 2 | let hits = vec![create_hit(0.9, "doc1"), create_hit(0.8, "doc2")]; |
828 | ||
829 | 2 | let shards = vec![create_shard_response(0, hits, 2)]; |
830 | ||
831 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
832 | ||
833 | 2 | assert_eq!(result.hits.len(), 2); |
834 | 2 | } |
835 | ||
836 | #[test] | |
837 | 2 | fn test_tie_breaking_by_primary_key() { |
838 | 2 | let merger = MergerImpl; |
839 | ||
840 | 2 | let hits = vec![ |
841 | 2 | create_hit(0.5, "zebra"), |
842 | 2 | create_hit(0.5, "apple"), |
843 | 2 | create_hit(0.5, "banana"), |
844 | ]; | |
845 | ||
846 | 2 | let shards = vec![create_shard_response(0, hits, 3)]; |
847 | ||
848 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
849 | ||
850 | // All have score 0.5, so should be sorted by primary_key ascending | |
851 | 2 | assert_eq!(result.hits[0]["id"], "apple"); |
852 | 2 | assert_eq!(result.hits[1]["id"], "banana"); |
853 | 2 | assert_eq!(result.hits[2]["id"], "zebra"); |
854 | 2 | } |
855 | ||
856 | #[test] | |
857 | 2 | fn test_stable_serialization_same_input_same_json() { |
858 | 2 | let shard1 = serde_json::json!({ |
859 | 2 | "hits": [], |
860 | 2 | "estimatedTotalHits": 0, |
861 | 2 | "processingTimeMs": 10, |
862 | 2 | "facetDistribution": { |
863 | 2 | "color": {"red": 10, "blue": 5}, |
864 | 2 | "size": {"large": 8} |
865 | } | |
866 | }); | |
867 | ||
868 | 2 | let shard2 = serde_json::json!({ |
869 | 2 | "hits": [], |
870 | 2 | "estimatedTotalHits": 0, |
871 | 2 | "processingTimeMs": 15, |
872 | 2 | "facetDistribution": { |
873 | 2 | "color": {"red": 7, "green": 3}, |
874 | 2 | "size": {"large": 4, "small": 6} |
875 | } | |
876 | }); | |
877 | ||
878 | 2 | let input = MergeInput { |
879 | 2 | shard_hits: vec![ |
880 | 2 | ShardHitPage { body: shard1, success: true }, |
881 | 2 | ShardHitPage { body: shard2, success: true }, |
882 | 2 | ], |
883 | 2 | offset: 0, |
884 | 2 | limit: 10, |
885 | 2 | client_requested_score: false, |
886 | 2 | facets: None, |
887 | 2 | }; |
888 | ||
889 | 2 | let result1 = merge(input.clone()); |
890 | 2 | let result2 = merge(input); |
891 | ||
892 | // Serialize both results to JSON | |
893 | 2 | let json1 = serde_json::to_string(&result1).unwrap(); |
894 | 2 | let json2 = serde_json::to_string(&result2).unwrap(); |
895 | ||
896 | // Must be byte-identical | |
897 | 2 | assert_eq!(json1, json2); |
898 | 2 | } |
899 | ||
900 | #[test] | |
901 | 2 | fn test_binary_heap_efficiency_large_fan_out() { |
902 | 2 | let merger = MergerImpl; |
903 | ||
904 | // Test that the heap correctly maintains top-k when we have more hits than heap_size | |
905 | 2 | let mut hits = vec![]; |
906 | 202 | for |
907 | 200 | hits.push(create_hit(i as f64 / 100.0, &format!("doc{}", i))); |
908 | 200 | } |
909 | ||
910 | 2 | let shards = vec![create_shard_response(0, hits, 100)]; |
911 | ||
912 | 2 | let result = merger.merge(shards, 0, 10, false).unwrap(); |
913 | ||
914 | // Should get the top 10 scores (0.90-0.99) | |
915 | 2 | assert_eq!(result.hits.len(), 10); |
916 | 2 | assert_eq!(result.hits[0]["id"], "doc99"); // score 0.99 |
917 | 2 | assert_eq!(result.hits[9]["id"], "doc90"); // score 0.90 |
918 | 2 | } |
919 | ||
920 | #[test] | |
921 | 2 | fn test_offset_limit_pagination_reconstruction() { |
922 | 2 | let merger = MergerImpl; |
923 | ||
924 | // Create 50 docs with known scores | |
925 | 2 | let mut hits = vec![]; |
926 | 102 | for |
927 | 100 | hits.push(create_hit((50 - i) as f64 / 100.0, &format!("doc{:02}", i))); |
928 | 100 | } |
929 | ||
930 | // Get all 50 in one go | |
931 | 2 | let shards_all = vec![create_shard_response(0, hits.clone(), 50)]; |
932 | 2 | let result_all = merger.merge(shards_all, 0, 50, false).unwrap(); |
933 | ||
934 | // Get 5 pages of 10 | |
935 | 2 | let mut paged_ids = vec![]; |
936 | 12 | for |
937 | 10 | let shards = vec![create_shard_response(0, hits.clone(), 50)]; |
938 | 10 | let result = merger.merge(shards, page * 10, 10, false).unwrap(); |
939 | 110 | for |
940 | 100 | paged_ids.push(hit["id"].as_str().unwrap().to_string()); |
941 | 100 | } |
942 | } | |
943 | ||
944 | // Concatenated pages should match the single limit=50 query | |
945 | 2 | let all_ids: Vec<_> = result_all |
946 | 2 | .hits |
947 | 2 | .iter() |
948 | 100 | . |
949 | 2 | .collect(); |
950 | ||
951 | 2 | assert_eq!(paged_ids, all_ids); |
952 | 2 | } |
953 | ||
954 | #[test] | |
955 | 2 | fn test_strip_all_miroir_reserved_fields() { |
956 | 2 | let merger = MergerImpl; |
957 | ||
958 | 2 | let hit = serde_json::json!({ |
959 | 2 | "id": "doc1", |
960 | 2 | "_rankingScore": 0.9, |
961 | 2 | "_miroir_shard": 0, |
962 | 2 | "_miroir_internal": "some_value", |
963 | 2 | "title": "Test", |
964 | }); | |
965 | ||
966 | 2 | let shards = vec![create_shard_response(0, vec![hit], 1)]; |
967 | ||
968 | 2 | let result = merger.merge(shards, 0, 10, true).unwrap(); |
969 | ||
970 | // _rankingScore should be present (requested) | |
971 | 2 | assert_eq!(result.hits[0]["_rankingScore"], 0.9); |
972 | // All _miroir_* fields should be stripped | |
973 | 2 | assert!(result.hits[0].get("_miroir_shard").is_none()); |
974 | 2 | assert!(result.hits[0].get("_miroir_internal").is_none()); |
975 | // Other fields should be present | |
976 | 2 | assert_eq!(result.hits[0]["title"], "Test"); |
977 | 2 | } |
978 | ||
979 | #[test] | |
980 | 2 | fn test_facet_filter_only_merges_requested_facets() { |
981 | 2 | let facets = serde_json::json!({ |
982 | 2 | "color": {"red": 10, "blue": 5}, |
983 | 2 | "size": {"large": 8} |
984 | }); | |
985 | ||
986 | 2 | let input = MergeInput { |
987 | 2 | shard_hits: vec![ShardHitPage { |
988 | 2 | body: serde_json::json!({ |
989 | 2 | "hits": [], |
990 | 2 | "estimatedTotalHits": 0, |
991 | 2 | "processingTimeMs": 10, |
992 | 2 | "facetDistribution": facets, |
993 | 2 | }), |
994 | 2 | success: true, |
995 | 2 | }], |
996 | 2 | offset: 0, |
997 | 2 | limit: 10, |
998 | 2 | client_requested_score: false, |
999 | 2 | facets: Some(vec!["color".to_string()]), |
1000 | 2 | }; |
1001 | ||
1002 | 2 | let result = merge(input); |
1003 | ||
1004 | // Only "color" should be present | |
1005 | 2 | assert!(result.facets.get("color").is_some()); |
1006 | 2 | assert!(result.facets.get("size").is_none()); |
1007 | 2 | } |
1008 | } |
Line | Count | Source |
1 | //! Shard migration cutover state machine. | |
2 | //! | |
3 | //! Implements the node-addition migration flow from plan §4 with explicit state | |
4 | //! transitions and a race-window-safe cutover sequence. | |
5 | //! | |
6 | //! ## Race window analysis (plan §15 OP#1) | |
7 | //! | |
8 | //! The dangerous window is between "mark node active" (routing changes to new-node-only) | |
9 | //! and "delete migrated shard from old node." A document written during dual-write that | |
10 | //! succeeded on OLD but failed on NEW — and arrived after the last migration page — | |
11 | //! would be deleted from OLD without ever reaching NEW. | |
12 | //! | |
13 | //! ## Solution: quiesce-then-verify cutover | |
14 | //! | |
15 | //! Instead of the naïve sequence (mark active → stop dual-write → delete old), we use: | |
16 | //! | |
17 | //! 1. Stop dual-write (no new writes go to either node for affected shards) | |
18 | //! 2. Drain: wait for all in-flight writes to both OLD and NEW to complete | |
19 | //! 3. Delta migration: re-read affected shards from OLD (catches anything written since | |
20 | //! the last migration page) and write deltas to NEW | |
21 | //! 4. Mark node active (routing switches to NEW-only) | |
22 | //! 5. Delete migrated shard from OLD | |
23 | //! | |
24 | //! Step 3 is the key: it closes the race window by ensuring NEW has a complete picture | |
25 | //! before we commit the routing change. The cost is one extra pagination pass over each | |
26 | //! migrated shard — bounded by the number of docs written during the migration window. | |
27 | ||
28 | use std::collections::{HashMap, HashSet}; | |
29 | use std::fmt; | |
30 | use std::time::{Duration, Instant}; | |
31 | ||
32 | use serde::{Deserialize, Serialize}; | |
33 | ||
34 | /// Unique identifier for a shard migration operation. | |
35 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] | |
36 | pub struct MigrationId(pub u64); | |
37 | ||
38 | impl fmt::Display for MigrationId { | |
39 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
40 | 0 | write!(f, "{}", self.0) |
41 | 0 | } |
42 | } | |
43 | ||
44 | /// Identifier for a physical node in the cluster. | |
45 | #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] | |
46 | pub struct NodeId(pub String); | |
47 | ||
48 | impl fmt::Display for NodeId { | |
49 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
50 | 0 | write!(f, "{}", self.0) |
51 | 0 | } |
52 | } | |
53 | ||
54 | /// Identifier for a logical shard. | |
55 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] | |
56 | pub struct ShardId(pub u32); | |
57 | ||
58 | impl fmt::Display for ShardId { | |
59 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
60 | 0 | write!(f, "s{}", self.0) |
61 | 0 | } |
62 | } | |
63 | ||
64 | /// Per-shard migration state within a node-addition migration. | |
65 | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] | |
66 | pub enum ShardMigrationState { | |
67 | /// Waiting for background migration to begin. | |
68 | Pending, | |
69 | /// Background pagination is reading docs from source and writing to target. | |
70 | Migrating { | |
71 | docs_copied: u64, | |
72 | pages_remaining: u32, | |
73 | }, | |
74 | /// Background migration complete, awaiting cutover. | |
75 | MigrationComplete { docs_copied: u64 }, | |
76 | /// Dual-write stopped, in-flight writes draining. | |
77 | Draining { | |
78 | in_flight_count: u32, | |
79 | docs_copied: u64, | |
80 | }, | |
81 | /// Delta pass: re-reading source to catch stragglers written during migration. | |
82 | DeltaPass { | |
83 | docs_copied: u64, | |
84 | delta_docs_copied: u64, | |
85 | }, | |
86 | /// Node is active for this shard; old replica data deleted. | |
87 | Active, | |
88 | /// Migration failed at this phase. | |
89 | Failed { phase: String, reason: String }, | |
90 | } | |
91 | ||
92 | impl fmt::Display for ShardMigrationState { | |
93 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
94 | 0 | match self { |
95 | 0 | Self::Pending => write!(f, "pending"), |
96 | Self::Migrating { | |
97 | 0 | docs_copied, |
98 | 0 | pages_remaining, |
99 | } => { | |
100 | 0 | write!( |
101 | 0 | f, |
102 | 0 | "migrating({docs_copied} copied, {pages_remaining} pages left)" |
103 | ) | |
104 | } | |
105 | 0 | Self::MigrationComplete { docs_copied } => { |
106 | 0 | write!(f, "migration_complete({docs_copied} copied)") |
107 | } | |
108 | Self::Draining { | |
109 | 0 | in_flight_count, |
110 | 0 | docs_copied, |
111 | } => { | |
112 | 0 | write!( |
113 | 0 | f, |
114 | 0 | "draining({in_flight_count} in-flight, {docs_copied} copied)" |
115 | ) | |
116 | } | |
117 | Self::DeltaPass { | |
118 | 0 | docs_copied, |
119 | 0 | delta_docs_copied, |
120 | } => { | |
121 | 0 | write!(f, "delta_pass({docs_copied} + {delta_docs_copied} copied)") |
122 | } | |
123 | 0 | Self::Active => write!(f, "active"), |
124 | 0 | Self::Failed { phase, reason } => write!(f, "failed({phase}: {reason})"), |
125 | } | |
126 | 0 | } |
127 | } | |
128 | ||
129 | /// Overall migration phase for a node addition. | |
130 | #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] | |
131 | pub enum MigrationPhase { | |
132 | /// Computing which shards move to the new node. | |
133 | ComputingAssignments, | |
134 | /// Dual-write active; background migration in progress. | |
135 | DualWriteMigrating, | |
136 | /// Background migration done; beginning cutover. | |
137 | CutoverBegin, | |
138 | /// Stopping dual-write; waiting for in-flight writes to settle. | |
139 | CutoverDraining, | |
140 | /// Re-reading source to catch docs written during migration. | |
141 | CutoverDeltaPass, | |
142 | /// Marking new node active; switching routing. | |
143 | CutoverActivate, | |
144 | /// Deleting migrated shard data from old nodes. | |
145 | CutoverCleanup, | |
146 | /// All shards migrated; migration complete. | |
147 | Complete, | |
148 | /// Migration failed. | |
149 | Failed(String), | |
150 | } | |
151 | ||
152 | impl fmt::Display for MigrationPhase { | |
153 | 0 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
154 | 0 | match self { |
155 | 0 | Self::ComputingAssignments => write!(f, "computing_assignments"), |
156 | 0 | Self::DualWriteMigrating => write!(f, "dual_write_migrating"), |
157 | 0 | Self::CutoverBegin => write!(f, "cutover_begin"), |
158 | 0 | Self::CutoverDraining => write!(f, "cutover_draining"), |
159 | 0 | Self::CutoverDeltaPass => write!(f, "cutover_delta_pass"), |
160 | 0 | Self::CutoverActivate => write!(f, "cutover_activate"), |
161 | 0 | Self::CutoverCleanup => write!(f, "cutover_cleanup"), |
162 | 0 | Self::Complete => write!(f, "complete"), |
163 | 0 | Self::Failed(msg) => write!(f, "failed({msg})"), |
164 | } | |
165 | 0 | } |
166 | } | |
167 | ||
168 | /// A single document write targeting a shard during migration. | |
169 | #[derive(Debug, Clone)] | |
170 | pub struct InFlightWrite { | |
171 | pub doc_id: String, | |
172 | pub shard: ShardId, | |
173 | pub target_nodes: Vec<NodeId>, | |
174 | pub completed_nodes: HashSet<NodeId>, | |
175 | pub failed_nodes: HashMap<NodeId, String>, | |
176 | pub submitted_at: Instant, | |
177 | } | |
178 | ||
179 | // Serialize Instant as a placeholder bool (present/absent). | |
180 | // Instant is monotonic and not meaningfully serializable across processes; | |
181 | // on deserialize, reconstruct as Instant::now(). | |
182 | mod instant_serde { | |
183 | use serde::{Deserialize, Deserializer, Serialize, Serializer}; | |
184 | use std::time::Instant; | |
185 | ||
186 | 0 | pub fn serialize<S>(instant: &Option<Instant>, serializer: S) -> Result<S::Ok, S::Error> |
187 | 0 | where |
188 | 0 | S: Serializer, |
189 | { | |
190 | 0 | instant.is_some().serialize(serializer) |
191 | 0 | } |
192 | ||
193 | 0 | pub fn deserialize<'de, D>(deserializer: D) -> Result<Option<Instant>, D::Error> |
194 | 0 | where |
195 | 0 | D: Deserializer<'de>, |
196 | { | |
197 | 0 | let present = bool::deserialize(deserializer)?; |
198 | 0 | Ok(if present { Some(Instant::now()) } else { None }) |
199 | 0 | } |
200 | } | |
201 | ||
202 | /// Configuration for migration cutover behavior. | |
203 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
204 | pub struct MigrationConfig { | |
205 | /// Maximum time to wait for in-flight writes to drain during cutover. | |
206 | pub drain_timeout: Duration, | |
207 | /// Whether to perform the delta pass (re-read source after stopping dual-write). | |
208 | /// Disabling this saves a pagination pass but opens the race window — only safe | |
209 | /// when anti-entropy is enabled as a safety net. | |
210 | pub skip_delta_pass: bool, | |
211 | /// Whether anti-entropy is enabled — used to determine if skip_delta_pass is safe. | |
212 | pub anti_entropy_enabled: bool, | |
213 | } | |
214 | ||
215 | impl Default for MigrationConfig { | |
216 | 16 | fn default() -> Self { |
217 | 16 | Self { |
218 | 16 | drain_timeout: Duration::from_secs(30), |
219 | 16 | skip_delta_pass: false, |
220 | 16 | anti_entropy_enabled: true, |
221 | 16 | } |
222 | 16 | } |
223 | } | |
224 | ||
225 | /// Error type for migration operations. | |
226 | #[derive(Debug, thiserror::Error)] | |
227 | pub enum MigrationError { | |
228 | #[error( | |
229 | "anti-entropy is disabled and delta pass is skipped — documents may be lost at cutover" | |
230 | )] | |
231 | UnsafeCutoverNoAntiEntropy, | |
232 | #[error("drain timeout exceeded: {0} in-flight writes still pending")] | |
233 | DrainTimeout(u32), | |
234 | #[error("shard {0} is not in a valid state for this transition (current: {1})")] | |
235 | InvalidTransition(ShardId, String), | |
236 | #[error("migration {0} not found")] | |
237 | NotFound(MigrationId), | |
238 | #[error("delta pass failed for shard {0}: {1}")] | |
239 | DeltaPassFailed(ShardId, String), | |
240 | } | |
241 | ||
242 | /// Tracks the state of a node-addition migration. | |
243 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
244 | pub struct MigrationState { | |
245 | pub id: MigrationId, | |
246 | pub new_node: NodeId, | |
247 | pub replica_group: u32, | |
248 | pub phase: MigrationPhase, | |
249 | pub affected_shards: HashMap<ShardId, ShardMigrationState>, | |
250 | /// Maps shard → old node that currently owns it. | |
251 | pub old_owners: HashMap<ShardId, NodeId>, | |
252 | #[serde(with = "instant_serde")] | |
253 | pub started_at: Option<Instant>, | |
254 | #[serde(with = "instant_serde")] | |
255 | pub completed_at: Option<Instant>, | |
256 | } | |
257 | ||
258 | /// The migration coordinator manages shard migration state transitions. | |
259 | pub struct MigrationCoordinator { | |
260 | config: MigrationConfig, | |
261 | migrations: HashMap<MigrationId, MigrationState>, | |
262 | next_id: u64, | |
263 | /// In-flight writes being tracked for drain during cutover. | |
264 | in_flight: Vec<InFlightWrite>, | |
265 | } | |
266 | ||
267 | impl MigrationCoordinator { | |
268 | 10 | pub fn new(config: MigrationConfig) -> Self { |
269 | 10 | Self { |
270 | 10 | config, |
271 | 10 | migrations: HashMap::new(), |
272 | 10 | next_id: 0, |
273 | 10 | in_flight: Vec::new(), |
274 | 10 | } |
275 | 10 | } |
276 | ||
277 | /// Validate migration safety before starting. Returns an error if the configuration | |
278 | /// would allow data loss at the cutover boundary. | |
279 | 10 | pub fn validate_safety(&self) -> Result<(), MigrationError> { |
280 | 10 | if self.config.skip_delta_pass && |
281 | 2 | return Err(MigrationError::UnsafeCutoverNoAntiEntropy); |
282 | 8 | } |
283 | 8 | Ok(()) |
284 | 10 | } |
285 | ||
286 | /// Begin a new node-addition migration. | |
287 | 10 | pub fn begin_migration( |
288 | 10 | &mut self, |
289 | 10 | new_node: NodeId, |
290 | 10 | replica_group: u32, |
291 | 10 | affected_shards: HashMap<ShardId, NodeId>, |
292 | 10 | ) -> Result<MigrationId, MigrationError> { |
293 | 10 | self.validate_safety() |
294 | ||
295 | 8 | let id = MigrationId(self.next_id); |
296 | 8 | self.next_id += 1; |
297 | ||
298 | 8 | let shard_states: HashMap<ShardId, ShardMigrationState> = affected_shards |
299 | 8 | .keys() |
300 | 10 | . |
301 | 8 | .collect(); |
302 | ||
303 | 8 | let state = MigrationState { |
304 | 8 | id, |
305 | 8 | new_node, |
306 | 8 | replica_group, |
307 | 8 | phase: MigrationPhase::ComputingAssignments, |
308 | 8 | affected_shards: shard_states, |
309 | 8 | old_owners: affected_shards, |
310 | 8 | started_at: Some(Instant::now()), |
311 | 8 | completed_at: None, |
312 | 8 | }; |
313 | ||
314 | 8 | self.migrations.insert(id, state); |
315 | 8 | Ok(id) |
316 | 10 | } |
317 | ||
318 | /// Transition to dual-write + background migration phase. | |
319 | 8 | pub fn begin_dual_write(&mut self, id: MigrationId) -> Result<(), MigrationError> { |
320 | 8 | let state = self |
321 | 8 | .migrations |
322 | 8 | .get_mut(&id) |
323 | 8 | .ok_or(MigrationError::NotFound(id)) |
324 | 8 | state.phase = MigrationPhase::DualWriteMigrating; |
325 | 10 | for shard_state in |
326 | 10 | if *shard_state == ShardMigrationState::Pending { |
327 | 10 | *shard_state = ShardMigrationState::Migrating { |
328 | 10 | docs_copied: 0, |
329 | 10 | pages_remaining: 0, |
330 | 10 | }; |
331 | 10 |
|
332 | } | |
333 | 8 | Ok(()) |
334 | 8 | } |
335 | ||
336 | /// Record that a shard's background migration completed. | |
337 | 10 | pub fn shard_migration_complete( |
338 | 10 | &mut self, |
339 | 10 | id: MigrationId, |
340 | 10 | shard: ShardId, |
341 | 10 | docs_copied: u64, |
342 | 10 | ) -> Result<(), MigrationError> { |
343 | 10 | let state = self |
344 | 10 | .migrations |
345 | 10 | .get_mut(&id) |
346 | 10 | .ok_or(MigrationError::NotFound(id)) |
347 | 10 | let shard_state = state.affected_shards.get_mut(&shard).ok_or_else(|| |
348 | 0 | MigrationError::InvalidTransition(shard, "shard not in migration".into()) |
349 | 0 | })?; |
350 | ||
351 | 10 | match shard_state { |
352 | 10 | ShardMigrationState::Migrating { .. } => { |
353 | 10 | *shard_state = ShardMigrationState::MigrationComplete { docs_copied }; |
354 | 10 | } |
355 | _ => { | |
356 | 0 | return Err(MigrationError::InvalidTransition( |
357 | 0 | shard, |
358 | 0 | shard_state.to_string(), |
359 | 0 | )); |
360 | } | |
361 | } | |
362 | ||
363 | // Check if all shards are done migrating | |
364 | 10 | let all_complete = state |
365 | 10 | .affected_shards |
366 | 10 | .values() |
367 | 13 | . |
368 | ||
369 | 10 | if all_complete { |
370 | 8 | state.phase = MigrationPhase::CutoverBegin; |
371 | 8 |
|
372 | ||
373 | 10 | Ok(()) |
374 | 10 | } |
375 | ||
376 | /// Begin the cutover sequence: stop dual-write and drain in-flight writes. | |
377 | 6 | pub fn begin_cutover(&mut self, id: MigrationId) -> Result<MigrationPhase, MigrationError> { |
378 | 6 | let state = self |
379 | 6 | .migrations |
380 | 6 | .get_mut(&id) |
381 | 6 | .ok_or(MigrationError::NotFound(id)) |
382 | ||
383 | 6 | if ! |
384 | 0 | return Err(MigrationError::InvalidTransition( |
385 | 0 | ShardId(0), |
386 | 0 | format!("expected CutoverBegin, got {}", state.phase), |
387 | 0 | )); |
388 | 6 | } |
389 | ||
390 | // Transition all shards to Draining | |
391 | 6 | let total_in_flight = self.in_flight.len() as u32; |
392 | 8 | for (shard, shard_state) in |
393 | 8 | match shard_state { |
394 | 8 | ShardMigrationState::MigrationComplete { docs_copied } => { |
395 | 8 | *shard_state = ShardMigrationState::Draining { |
396 | 8 | in_flight_count: total_in_flight, |
397 | 8 | docs_copied: *docs_copied, |
398 | 8 | }; |
399 | 8 | } |
400 | _ => { | |
401 | 0 | return Err(MigrationError::InvalidTransition( |
402 | 0 | *shard, |
403 | 0 | shard_state.to_string(), |
404 | 0 | )); |
405 | } | |
406 | } | |
407 | } | |
408 | ||
409 | 6 | state.phase = MigrationPhase::CutoverDraining; |
410 | 6 | Ok(state.phase.clone()) |
411 | 6 | } |
412 | ||
413 | /// Register an in-flight write for tracking during drain. | |
414 | 4 | pub fn register_in_flight(&mut self, write: InFlightWrite) { |
415 | 4 | self.in_flight.push(write); |
416 | 4 | } |
417 | ||
418 | /// Acknowledge completion of a write to a specific node. | |
419 | 0 | pub fn ack_write(&mut self, doc_id: &str, node: &NodeId) { |
420 | 0 | for write in &mut self.in_flight { |
421 | 0 | if write.doc_id == doc_id { |
422 | 0 | write.completed_nodes.insert(node.clone()); |
423 | 0 | } |
424 | } | |
425 | 0 | } |
426 | ||
427 | /// Mark a write as failed on a specific node. | |
428 | 0 | pub fn fail_write(&mut self, doc_id: &str, node: &NodeId, reason: String) { |
429 | 0 | for write in &mut self.in_flight { |
430 | 0 | if write.doc_id == doc_id { |
431 | 0 | write.failed_nodes.insert(node.clone(), reason.clone()); |
432 | 0 | } |
433 | } | |
434 | 0 | } |
435 | ||
436 | /// Check if all in-flight writes have completed (drained). | |
437 | 6 | pub fn is_drained(&self) -> bool { |
438 | 6 | self.in_flight |
439 | 6 | .iter() |
440 | 6 | .all(|w| |
441 | 6 | } |
442 | ||
443 | /// Complete the drain and move to delta pass or activation. | |
444 | 6 | pub fn complete_drain(&mut self, id: MigrationId) -> Result<MigrationPhase, MigrationError> { |
445 | // First check phase exists without holding mutable borrow | |
446 | 6 | let phase = self |
447 | 6 | .migrations |
448 | 6 | .get(&id) |
449 | 6 | .ok_or(MigrationError::NotFound(id)) |
450 | .phase | |
451 | 6 | .clone(); |
452 | ||
453 | 6 | if ! |
454 | 0 | return Err(MigrationError::InvalidTransition( |
455 | 0 | ShardId(0), |
456 | 0 | format!("expected CutoverDraining, got {phase}"), |
457 | 0 | )); |
458 | 6 | } |
459 | ||
460 | // Check drain status | |
461 | 6 | if !self.is_drained() { |
462 | 2 | let remaining = self |
463 | 2 | .in_flight |
464 | 2 | .iter() |
465 | 2 | .filter(|w| w.completed_nodes.len() + w.failed_nodes.len() < w.target_nodes.len()) |
466 | 2 | .count() as u32; |
467 | 2 | return Err(MigrationError::DrainTimeout(remaining)); |
468 | 4 | } |
469 | ||
470 | // Collect docs that need delta pass | |
471 | 4 | let needs_delta = self.collect_delta_candidates(id) |
472 | 4 | let skip_delta = self.config.skip_delta_pass; |
473 | ||
474 | // Now get mutable borrow to update state | |
475 | 4 | let state = self |
476 | 4 | .migrations |
477 | 4 | .get_mut(&id) |
478 | 4 | .ok_or(MigrationError::NotFound(id)) |
479 | ||
480 | 4 | if skip_delta { |
481 | 2 | // Skip delta pass — safe only if anti-entropy is enabled |
482 | 2 | state.phase = MigrationPhase::CutoverActivate; |
483 | 2 | } else if needs_delta.is_empty() { |
484 | 0 | state.phase = MigrationPhase::CutoverActivate; |
485 | 0 | } else { |
486 | 2 | state.phase = MigrationPhase::CutoverDeltaPass; |
487 | 4 | for (_shard, shard_state) in |
488 | 4 | if let ShardMigrationState::Draining { docs_copied, .. } = shard_state { |
489 | 4 | *shard_state = ShardMigrationState::DeltaPass { |
490 | 4 | docs_copied: *docs_copied, |
491 | 4 | delta_docs_copied: 0, |
492 | 4 | }; |
493 | 4 |
|
494 | } | |
495 | } | |
496 | ||
497 | // Clear only the in-flight writes for this migration | |
498 | 4 | let affected_shards = state |
499 | 4 | .affected_shards |
500 | 4 | .keys() |
501 | 4 | .cloned() |
502 | 4 | .collect::<HashSet<_>>(); |
503 | 4 | self.in_flight |
504 | 4 | .retain(|w| ! |
505 | ||
506 | // If going to activate, do that now (drop mutable borrow first) | |
507 | 4 | let next_phase = state.phase.clone(); |
508 | 4 | if |
509 | 2 | let _ = state; |
510 | 2 | self.activate_shards(id) |
511 | // Return the new phase after activation | |
512 | 2 | return Ok(self |
513 | 2 | .migrations |
514 | 2 | .get(&id) |
515 | 2 | .map(|s| s.phase.clone()) |
516 | 2 | .unwrap_or(MigrationPhase::CutoverCleanup)); |
517 | 2 | } |
518 | ||
519 | 2 | Ok(next_phase) |
520 | 6 | } |
521 | ||
522 | /// Identify writes that need the delta pass — those that succeeded on OLD but | |
523 | /// failed (or never reached) NEW. | |
524 | 4 | fn collect_delta_candidates( |
525 | 4 | &self, |
526 | 4 | id: MigrationId, |
527 | 4 | ) -> Result<HashMap<ShardId, Vec<String>>, MigrationError> { |
528 | 4 | let state = self |
529 | 4 | .migrations |
530 | 4 | .get(&id) |
531 | 4 | .ok_or(MigrationError::NotFound(id)) |
532 | 4 | let mut candidates: HashMap<ShardId, Vec<String>> = HashMap::new(); |
533 | ||
534 | 6 | for |
535 | 2 | let old_owner = match state.old_owners.get(&write.shard) { |
536 | 2 | Some(owner) => owner, |
537 | 0 | None => continue, |
538 | }; | |
539 | ||
540 | 2 | let succeeded_on_old = write.completed_nodes.contains(old_owner); |
541 | 2 | let succeeded_on_new = write.completed_nodes.contains(&state.new_node); |
542 | ||
543 | // Doc is on OLD but not on NEW — delta pass must catch it | |
544 | 2 | if succeeded_on_old && !succeeded_on_new { |
545 | 2 | candidates |
546 | 2 | .entry(write.shard) |
547 | 2 | .or_default() |
548 | 2 | .push(write.doc_id.clone()); |
549 | 2 |
|
550 | } | |
551 | ||
552 | 4 | Ok(candidates) |
553 | 4 | } |
554 | ||
555 | /// Record that the delta pass completed for a shard. | |
556 | 4 | pub fn shard_delta_complete( |
557 | 4 | &mut self, |
558 | 4 | id: MigrationId, |
559 | 4 | shard: ShardId, |
560 | 4 | delta_docs: u64, |
561 | 4 | ) -> Result<(), MigrationError> { |
562 | 4 | let state = self |
563 | 4 | .migrations |
564 | 4 | .get_mut(&id) |
565 | 4 | .ok_or(MigrationError::NotFound(id)) |
566 | 4 | let shard_state = state.affected_shards.get_mut(&shard).ok_or_else(|| |
567 | 0 | MigrationError::InvalidTransition(shard, "shard not in migration".into()) |
568 | 0 | })?; |
569 | ||
570 | 4 | match shard_state { |
571 | 4 | ShardMigrationState::DeltaPass { docs_copied, .. } => { |
572 | 4 | *shard_state = ShardMigrationState::MigrationComplete { |
573 | 4 | docs_copied: *docs_copied + delta_docs, |
574 | 4 | }; |
575 | 4 | } |
576 | _ => { | |
577 | 0 | return Err(MigrationError::InvalidTransition( |
578 | 0 | shard, |
579 | 0 | shard_state.to_string(), |
580 | 0 | )); |
581 | } | |
582 | } | |
583 | ||
584 | // Check if all shards done with delta | |
585 | 4 | let all_complete = state |
586 | 4 | .affected_shards |
587 | 4 | .values() |
588 | 7 | . |
589 | ||
590 | 4 | if all_complete { |
591 | 2 | state.phase = MigrationPhase::CutoverActivate; |
592 | 2 | self.activate_shards(id) |
593 | 2 | } |
594 | ||
595 | 4 | Ok(()) |
596 | 4 | } |
597 | ||
598 | /// Mark all affected shards as active on the new node. | |
599 | 4 | fn activate_shards(&mut self, id: MigrationId) -> Result<(), MigrationError> { |
600 | 4 | let state = self |
601 | 4 | .migrations |
602 | 4 | .get_mut(&id) |
603 | 4 | .ok_or(MigrationError::NotFound(id)) |
604 | ||
605 | 6 | for shard_state in |
606 | 6 | match shard_state { |
607 | ShardMigrationState::MigrationComplete { .. } | |
608 | 6 | | ShardMigrationState::Draining { .. } => { |
609 | 6 | *shard_state = ShardMigrationState::Active; |
610 | 6 | } |
611 | 0 | _ => {} |
612 | } | |
613 | } | |
614 | ||
615 | 4 | if |
616 | 4 | state.phase = MigrationPhase::CutoverCleanup; |
617 | 4 |
|
618 | ||
619 | 4 | Ok(()) |
620 | 4 | } |
621 | ||
622 | /// Complete the migration by deleting migrated shard data from old nodes. | |
623 | 4 | pub fn complete_cleanup(&mut self, id: MigrationId) -> Result<(), MigrationError> { |
624 | 4 | let state = self |
625 | 4 | .migrations |
626 | 4 | .get_mut(&id) |
627 | 4 | .ok_or(MigrationError::NotFound(id)) |
628 | ||
629 | 4 | if ! |
630 | 0 | return Err(MigrationError::InvalidTransition( |
631 | 0 | ShardId(0), |
632 | 0 | format!("expected CutoverCleanup, got {}", state.phase), |
633 | 0 | )); |
634 | 4 | } |
635 | ||
636 | 4 | state.phase = MigrationPhase::Complete; |
637 | 4 | state.completed_at = Some(Instant::now()); |
638 | 4 | Ok(()) |
639 | 4 | } |
640 | ||
641 | /// Get the current state of a migration. | |
642 | 6 | pub fn get_state(&self, id: MigrationId) -> Option<&MigrationState> { |
643 | 6 | self.migrations.get(&id) |
644 | 6 | } |
645 | ||
646 | /// Check if a write should go to both old and new node (dual-write phase). | |
647 | 6 | pub fn is_dual_write_active(&self, shard: ShardId) -> bool { |
648 | 6 | self.migrations.values().any(|m| { |
649 | 6 |
|
650 | 2 | && matches!( |
651 | 4 | m.affected_shards.get(&shard), |
652 | Some(ShardMigrationState::Migrating { .. }) | |
653 | ) | |
654 | 6 | }) |
655 | 6 | } |
656 | ||
657 | /// Get the migration config. | |
658 | 0 | pub fn config(&self) -> &MigrationConfig { |
659 | 0 | &self.config |
660 | 0 | } |
661 | } | |
662 | ||
663 | #[cfg(test)] | |
664 | mod tests { | |
665 | use super::*; | |
666 | ||
667 | 34 | fn node(s: &str) -> NodeId { |
668 | 34 | NodeId(s.to_string()) |
669 | 34 | } |
670 | ||
671 | 36 | fn shard(id: u32) -> ShardId { |
672 | 36 | ShardId(id) |
673 | 36 | } |
674 | ||
675 | #[test] | |
676 | 2 | fn test_safe_cutover_with_delta_pass() { |
677 | 2 | let config = MigrationConfig { |
678 | 2 | anti_entropy_enabled: false, |
679 | 2 | skip_delta_pass: false, |
680 | 2 | ..Default::default() |
681 | 2 | }; |
682 | 2 | let mut coord = MigrationCoordinator::new(config); |
683 | ||
684 | 2 | let affected = HashMap::from([(shard(0), node("old-0")), (shard(1), node("old-0"))]); |
685 | ||
686 | 2 | let mid = coord.begin_migration(node("new-0"), 0, affected).unwrap(); |
687 | 2 | coord.begin_dual_write(mid).unwrap(); |
688 | ||
689 | // Simulate background migration completing | |
690 | 2 | coord.shard_migration_complete(mid, shard(0), 500).unwrap(); |
691 | 2 | coord.shard_migration_complete(mid, shard(1), 300).unwrap(); |
692 | ||
693 | // Register an in-flight write that succeeded on OLD but not NEW. | |
694 | // The write must be marked as failed on NEW so is_drained() sees | |
695 | // completed + failed == target count. | |
696 | 2 | coord.register_in_flight(InFlightWrite { |
697 | 2 | doc_id: "doc-at-boundary".into(), |
698 | 2 | shard: shard(0), |
699 | 2 | target_nodes: vec![node("old-0"), node("new-0")], |
700 | 2 | completed_nodes: HashSet::from([node("old-0")]), |
701 | 2 | failed_nodes: HashMap::from([(node("new-0"), "write failed".into())]), |
702 | 2 | submitted_at: Instant::now(), |
703 | 2 | }); |
704 | ||
705 | // Cutover | |
706 | 2 | coord.begin_cutover(mid).unwrap(); |
707 | ||
708 | // The drain sees the in-flight write completed (on old, not on new) | |
709 | // Delta pass should be triggered | |
710 | 2 | let phase = coord.complete_drain(mid).unwrap(); |
711 | 2 | assert_eq!(phase, MigrationPhase::CutoverDeltaPass); |
712 | ||
713 | // Delta pass catches the straggler | |
714 | 2 | coord.shard_delta_complete(mid, shard(0), 1).unwrap(); |
715 | // Shard 1 had no stragglers, but needs delta complete too | |
716 | 2 | coord.shard_delta_complete(mid, shard(1), 0).unwrap(); |
717 | ||
718 | // Now activation and cleanup | |
719 | 2 | let state = coord.get_state(mid).unwrap(); |
720 | 2 | assert_eq!(state.phase, MigrationPhase::CutoverCleanup); |
721 | ||
722 | 2 | coord.complete_cleanup(mid).unwrap(); |
723 | 2 | let state = coord.get_state(mid).unwrap(); |
724 | 2 | assert_eq!(state.phase, MigrationPhase::Complete); |
725 | 2 | } |
726 | ||
727 | #[test] | |
728 | 2 | fn test_unsafe_cutover_refused_without_anti_entropy() { |
729 | 2 | let config = MigrationConfig { |
730 | 2 | anti_entropy_enabled: false, |
731 | 2 | skip_delta_pass: true, |
732 | 2 | ..Default::default() |
733 | 2 | }; |
734 | 2 | let mut coord = MigrationCoordinator::new(config); |
735 | ||
736 | 2 | let affected = HashMap::from([(shard(0), node("old-0"))]); |
737 | 2 | let result = coord.begin_migration(node("new-0"), 0, affected); |
738 | ||
739 | 2 | assert!(result.is_err()); |
740 | 2 | let err = result.unwrap_err(); |
741 | 2 | assert!( |
742 | 2 | } |
743 | ||
744 | #[test] | |
745 | 2 | fn test_skip_delta_pass_allowed_with_anti_entropy() { |
746 | 2 | let config = MigrationConfig { |
747 | 2 | anti_entropy_enabled: true, |
748 | 2 | skip_delta_pass: true, |
749 | 2 | ..Default::default() |
750 | 2 | }; |
751 | 2 | let mut coord = MigrationCoordinator::new(config); |
752 | ||
753 | 2 | let affected = HashMap::from([(shard(0), node("old-0"))]); |
754 | 2 | let mid = coord.begin_migration(node("new-0"), 0, affected).unwrap(); |
755 | 2 | coord.begin_dual_write(mid).unwrap(); |
756 | 2 | coord.shard_migration_complete(mid, shard(0), 100).unwrap(); |
757 | ||
758 | 2 | coord.begin_cutover(mid).unwrap(); |
759 | ||
760 | // With skip_delta_pass=true and AE enabled, drain goes straight to activate | |
761 | 2 | let phase = coord.complete_drain(mid).unwrap(); |
762 | 2 | assert_eq!(phase, MigrationPhase::CutoverCleanup); |
763 | ||
764 | 2 | coord.complete_cleanup(mid).unwrap(); |
765 | 2 | assert_eq!( |
766 | 2 | coord.get_state(mid).unwrap().phase, |
767 | MigrationPhase::Complete | |
768 | ); | |
769 | 2 | } |
770 | ||
771 | #[test] | |
772 | 2 | fn test_drain_timeout_blocks_cutover() { |
773 | 2 | let config = MigrationConfig { |
774 | 2 | anti_entropy_enabled: true, |
775 | 2 | skip_delta_pass: true, |
776 | 2 | ..Default::default() |
777 | 2 | }; |
778 | 2 | let mut coord = MigrationCoordinator::new(config); |
779 | ||
780 | 2 | let affected = HashMap::from([(shard(0), node("old-0"))]); |
781 | 2 | let mid = coord.begin_migration(node("new-0"), 0, affected).unwrap(); |
782 | 2 | coord.begin_dual_write(mid).unwrap(); |
783 | 2 | coord.shard_migration_complete(mid, shard(0), 100).unwrap(); |
784 | 2 | coord.begin_cutover(mid).unwrap(); |
785 | ||
786 | // Register an in-flight write that hasn't completed on either node | |
787 | 2 | coord.register_in_flight(InFlightWrite { |
788 | 2 | doc_id: "stuck-doc".into(), |
789 | 2 | shard: shard(0), |
790 | 2 | target_nodes: vec![node("old-0"), node("new-0")], |
791 | 2 | completed_nodes: HashSet::new(), |
792 | 2 | failed_nodes: HashMap::new(), |
793 | 2 | submitted_at: Instant::now(), |
794 | 2 | }); |
795 | ||
796 | // Drain should fail — write still in flight | |
797 | 2 | let result = coord.complete_drain(mid); |
798 | 2 | assert!(result.is_err()); |
799 | 2 | assert!( |
800 | 2 | result.unwrap_err(), |
801 | MigrationError::DrainTimeout(1) | |
802 | )); | |
803 | 2 | } |
804 | ||
805 | #[test] | |
806 | 2 | fn test_dual_write_tracking() { |
807 | 2 | let config = MigrationConfig::default(); |
808 | 2 | let mut coord = MigrationCoordinator::new(config); |
809 | ||
810 | 2 | let affected = HashMap::from([(shard(5), node("old-0"))]); |
811 | 2 | let mid = coord.begin_migration(node("new-0"), 0, affected).unwrap(); |
812 | 2 | coord.begin_dual_write(mid).unwrap(); |
813 | ||
814 | // Shard 5 is in dual-write | |
815 | 2 | assert!(coord.is_dual_write_active(shard(5))); |
816 | // Shard 99 is not being migrated | |
817 | 2 | assert!(!coord.is_dual_write_active(shard(99))); |
818 | ||
819 | // After migration completes, shard 5 is no longer dual-write | |
820 | 2 | coord.shard_migration_complete(mid, shard(5), 100).unwrap(); |
821 | 2 | assert!(!coord.is_dual_write_active(shard(5))); |
822 | 2 | } |
823 | } |
Line | Count | Source |
1 | //! Online resharding: window guard, simulation model, and load estimation. | |
2 | //! | |
3 | //! Implements the plan §13.1 shadow-index resharding mechanics and §15 OP#3 | |
4 | //! empirical validation of the 2× transient load caveat. | |
5 | ||
6 | use crate::router::{assign_shard_in_group, shard_for_key}; | |
7 | use crate::topology::{Group, NodeId}; | |
8 | use serde::{Deserialize, Serialize}; | |
9 | use std::time::SystemTime; | |
10 | ||
11 | // --------------------------------------------------------------------------- | |
12 | // Schedule window guard | |
13 | // --------------------------------------------------------------------------- | |
14 | ||
15 | /// A UTC time window like `"02:00-06:00"`. | |
16 | #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] | |
17 | pub struct TimeWindow { | |
18 | /// Start hour+minute in minutes since midnight UTC. | |
19 | pub start_mins: u16, | |
20 | /// End hour+minute in minutes since midnight UTC. | |
21 | pub end_mins: u16, | |
22 | } | |
23 | ||
24 | impl std::fmt::Display for TimeWindow { | |
25 | 0 | fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { |
26 | 0 | write!( |
27 | 0 | f, |
28 | 0 | "{:02}:{:02}-{:02}:{:02}", |
29 | 0 | self.start_mins / 60, |
30 | 0 | self.start_mins % 60, |
31 | 0 | self.end_mins / 60, |
32 | 0 | self.end_mins % 60 |
33 | ) | |
34 | 0 | } |
35 | } | |
36 | ||
37 | impl TimeWindow { | |
38 | /// Parse a `"HH:MM-HH:MM"` string (UTC). | |
39 | 32 | pub fn parse(s: &str) -> Result<Self, String> { |
40 | 32 | let (start, end) = s |
41 | 32 | .split_once('-') |
42 | 32 | .ok_or_else(|| format!( |
43 | Ok(TimeWindow { | |
44 | 32 | start_mins: Self::parse_hm(start) |
45 | 26 | end_mins: Self::parse_hm(end) |
46 | }) | |
47 | 32 | } |
48 | ||
49 | 58 | fn parse_hm(hm: &str) -> Result<u16, String> { |
50 | 58 | let ( |
51 | 58 | .split_once(':') |
52 | 58 | .ok_or_else(|| format!( |
53 | 56 | let h: u16 = h.parse().map_err(|_| format!( |
54 | 56 | let m: u16 = m.parse().map_err(|_| format!( |
55 | 56 | if h >= 24 || |
56 | 4 | return Err(format!("time out of range: {hm}")); |
57 | 52 | } |
58 | 52 | Ok(h * 60 + m) |
59 | 58 | } |
60 | ||
61 | /// Does `utc_minutes` (minutes since midnight UTC) fall inside this window? | |
62 | 30 | pub fn contains(&self, utc_minutes: u16) -> bool { |
63 | 30 | if self.start_mins <= self.end_mins { |
64 | 24 | utc_minutes >= self.start_mins && |
65 | } else { | |
66 | // Wraps midnight, e.g. 22:00-06:00 | |
67 | 6 | utc_minutes >= self.start_mins || |
68 | } | |
69 | 30 | } |
70 | } | |
71 | ||
72 | /// Resharding configuration (plan §13.1 + schedule window guard). | |
73 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
74 | pub struct ReshardingConfig { | |
75 | #[serde(default = "default_true")] | |
76 | pub enabled: bool, | |
77 | #[serde(default = "default_backfill_concurrency")] | |
78 | pub backfill_concurrency: usize, | |
79 | #[serde(default = "default_backfill_batch_size")] | |
80 | pub backfill_batch_size: usize, | |
81 | #[serde(default)] | |
82 | pub throttle_docs_per_sec: u64, | |
83 | #[serde(default = "default_true")] | |
84 | pub verify_before_swap: bool, | |
85 | #[serde(default = "default_retain_hours")] | |
86 | pub retain_old_index_hours: u64, | |
87 | /// Allowed schedule windows in `"HH:MM-HH:MM UTC"` format. | |
88 | /// Empty means any time is allowed (no restriction). | |
89 | #[serde(default)] | |
90 | pub allowed_windows: Vec<String>, | |
91 | } | |
92 | ||
93 | 8 | fn default_backfill_concurrency() -> usize { |
94 | 8 | 4 |
95 | 8 | } |
96 | 8 | fn default_backfill_batch_size() -> usize { |
97 | 8 | 1000 |
98 | 8 | } |
99 | 0 | fn default_true() -> bool { |
100 | 0 | true |
101 | 0 | } |
102 | 8 | fn default_retain_hours() -> u64 { |
103 | 8 | 48 |
104 | 8 | } |
105 | ||
106 | impl Default for ReshardingConfig { | |
107 | 8 | fn default() -> Self { |
108 | 8 | Self { |
109 | 8 | enabled: true, |
110 | 8 | backfill_concurrency: default_backfill_concurrency(), |
111 | 8 | backfill_batch_size: default_backfill_batch_size(), |
112 | 8 | throttle_docs_per_sec: 0, |
113 | 8 | verify_before_swap: true, |
114 | 8 | retain_old_index_hours: default_retain_hours(), |
115 | 8 | allowed_windows: Vec::new(), |
116 | 8 | } |
117 | 8 | } |
118 | } | |
119 | ||
120 | /// Result of the schedule window guard check. | |
121 | #[derive(Debug, Clone, PartialEq, Eq)] | |
122 | pub enum WindowGuardResult { | |
123 | /// Current time is inside an allowed window. | |
124 | Allowed { window: String }, | |
125 | /// No windows configured — always allowed. | |
126 | NoRestriction, | |
127 | /// Current time is outside all allowed windows. | |
128 | Denied { | |
129 | utc_now: String, | |
130 | allowed: Vec<String>, | |
131 | }, | |
132 | } | |
133 | ||
134 | /// Check whether resharding is allowed at the given UTC minute-of-day. | |
135 | /// | |
136 | /// Returns `Allowed` if `utc_minute` falls inside any configured window, | |
137 | /// `NoRestriction` if no windows are configured, or `Denied` otherwise. | |
138 | 12 | pub fn check_window(utc_minute: u16, config: &ReshardingConfig) -> WindowGuardResult { |
139 | 12 | if config.allowed_windows.is_empty() { |
140 | 2 | return WindowGuardResult::NoRestriction; |
141 | 10 | } |
142 | ||
143 | 18 | for |
144 | 14 | let window = match TimeWindow::parse(raw) { |
145 | 14 | Ok(w) => w, |
146 | 0 | Err(_) => continue, |
147 | }; | |
148 | 14 | if window.contains(utc_minute) { |
149 | 6 | return WindowGuardResult::Allowed { |
150 | 6 | window: raw.clone(), |
151 | 6 | }; |
152 | 8 | } |
153 | } | |
154 | ||
155 | 4 | WindowGuardResult::Denied { |
156 | 4 | utc_now: format!("{:02}:{:02} UTC", utc_minute / 60, utc_minute % 60), |
157 | 4 | allowed: config.allowed_windows.clone(), |
158 | 4 | } |
159 | 12 | } |
160 | ||
161 | /// Check the schedule window against the system clock. | |
162 | 0 | pub fn check_window_now(config: &ReshardingConfig) -> WindowGuardResult { |
163 | 0 | let utc_minute = current_utc_minute(); |
164 | 0 | check_window(utc_minute, config) |
165 | 0 | } |
166 | ||
167 | 0 | fn current_utc_minute() -> u16 { |
168 | 0 | let duration = SystemTime::now() |
169 | 0 | .duration_since(SystemTime::UNIX_EPOCH) |
170 | 0 | .unwrap_or_default(); |
171 | 0 | ((duration.as_secs() / 60) % (24 * 60)) as u16 |
172 | 0 | } |
173 | ||
174 | // --------------------------------------------------------------------------- | |
175 | // Resharding load simulation | |
176 | // --------------------------------------------------------------------------- | |
177 | ||
178 | /// Parameters for a single simulation run. | |
179 | #[derive(Debug, Clone)] | |
180 | pub struct SimParams { | |
181 | /// Document size in bytes. | |
182 | pub doc_size_bytes: u64, | |
183 | /// Total corpus size in bytes. | |
184 | pub corpus_size_bytes: u64, | |
185 | /// Incoming write rate in documents per second. | |
186 | pub write_rate_dps: u64, | |
187 | /// Number of replica groups. | |
188 | pub replica_groups: u32, | |
189 | /// Replication factor (intra-group copies per shard). | |
190 | pub replication_factor: usize, | |
191 | /// Old shard count (before reshard). | |
192 | pub old_shards: u32, | |
193 | /// New shard count (after reshard). | |
194 | pub new_shards: u32, | |
195 | /// Number of nodes per replica group. | |
196 | pub nodes_per_group: usize, | |
197 | /// Backfill throttle in documents per second (0 = unlimited). | |
198 | pub backfill_throttle_dps: u64, | |
199 | } | |
200 | ||
201 | /// Results from a single simulation run. | |
202 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
203 | pub struct SimResult { | |
204 | pub label: String, | |
205 | pub doc_size_bytes: u64, | |
206 | pub corpus_size_bytes: u64, | |
207 | pub total_docs: u64, | |
208 | pub replica_groups: u32, | |
209 | pub replication_factor: usize, | |
210 | pub old_shards: u32, | |
211 | pub new_shards: u32, | |
212 | pub nodes_per_group: usize, | |
213 | pub write_rate_dps: u64, | |
214 | ||
215 | /// Normal steady-state storage across entire cluster (bytes). | |
216 | pub normal_storage_bytes: u64, | |
217 | /// Peak storage during resharding (live + shadow full, bytes). | |
218 | pub peak_storage_bytes: u64, | |
219 | /// Storage amplification factor (peak / normal). | |
220 | pub storage_amplification: f64, | |
221 | ||
222 | /// Normal steady-state write rate (actual node writes/sec). | |
223 | pub normal_write_rate: u64, | |
224 | /// Dual-write rate (phase 2 only, no backfill). | |
225 | pub dual_write_rate: u64, | |
226 | /// Peak write rate during backfill + dual-write. | |
227 | pub peak_write_rate: u64, | |
228 | /// Write amplification factor during dual-write only. | |
229 | pub dual_write_amplification: f64, | |
230 | /// Write amplification factor during peak (backfill + dual-write). | |
231 | pub peak_write_amplification: f64, | |
232 | ||
233 | /// Backfill duration in seconds (at configured throttle). | |
234 | pub backfill_duration_secs: f64, | |
235 | /// Total bytes written during full reshard operation. | |
236 | pub total_bytes_written: u64, | |
237 | ||
238 | /// Per-node peak storage (bytes). | |
239 | pub per_node_peak_storage_bytes: u64, | |
240 | /// Per-node normal storage (bytes). | |
241 | pub per_node_normal_storage_bytes: u64, | |
242 | ||
243 | /// Hash distribution stats for old shards. | |
244 | pub old_shard_cv: f64, | |
245 | /// Hash distribution stats for new shards. | |
246 | pub new_shard_cv: f64, | |
247 | } | |
248 | ||
249 | /// Run a resharding load simulation with the given parameters. | |
250 | /// | |
251 | /// This models the six-phase resharding process from plan §13.1 using | |
252 | /// the actual routing code to compute shard assignments and estimate | |
253 | /// storage/write load. Document keys are synthesized for the corpus size | |
254 | /// and routed through the real hash function to measure distribution. | |
255 | 6 | pub fn simulate(params: &SimParams) -> SimResult { |
256 | 6 | let total_docs = params.corpus_size_bytes / params.doc_size_bytes; |
257 | 6 | let rf = params.replication_factor; |
258 | 6 | let rg = params.replica_groups; |
259 | 6 | let nodes_per_group = params.nodes_per_group; |
260 | ||
261 | // Build a synthetic topology for the simulation. | |
262 | 6 | let groups: Vec<Group> = (0..rg) |
263 | 10 | . |
264 | 10 | let mut group = Group::new(g); |
265 | 32 | for n in 0.. |
266 | 32 | group.add_node(NodeId::new(format!("node-g{g}-n{n}"))); |
267 | 32 | } |
268 | 10 | group |
269 | 10 | }) |
270 | 6 | .collect(); |
271 | ||
272 | // Simulate document distribution across old and new shard counts. | |
273 | // Use the actual router hash to get realistic distribution. | |
274 | 6 | let mut old_shard_counts: Vec<u64> = vec![0; params.old_shards as usize]; |
275 | 6 | let mut new_shard_counts: Vec<u64> = vec![0; params.new_shards as usize]; |
276 | ||
277 | // Track per-node storage for old and new shard assignments. | |
278 | // Each group stores the full corpus; each node in a group stores its | |
279 | // rendezvous-assigned fraction. | |
280 | 6 | let total_nodes = (rg as usize) * nodes_per_group; |
281 | 6 | let mut node_storage_old: Vec<u64> = vec![0; total_nodes]; |
282 | 6 | let mut node_storage_new: Vec<u64> = vec![0; total_nodes]; |
283 | ||
284 | 22.9M | for i in 0.. |
285 | 22.9M | let key = format!("doc-{i}"); |
286 | 22.9M | let old_shard = shard_for_key(&key, params.old_shards); |
287 | 22.9M | let new_shard = shard_for_key(&key, params.new_shards); |
288 | ||
289 | 22.9M | old_shard_counts[old_shard as usize] += 1; |
290 | 22.9M | new_shard_counts[new_shard as usize] += 1; |
291 | ||
292 | // For each replica group, assign shard to RF nodes. | |
293 | 43.9M | for (g_idx, group) in |
294 | 43.9M | let old_targets = assign_shard_in_group(old_shard, group.nodes(), rf); |
295 | 43.9M | let new_targets = assign_shard_in_group(new_shard, group.nodes(), rf); |
296 | ||
297 | 87.8M | for |
298 | 43.9M | let node_idx = g_idx * nodes_per_group |
299 | 91.5M | + |
300 | 43.9M | node_storage_old[node_idx] += params.doc_size_bytes; |
301 | } | |
302 | 87.8M | for |
303 | 43.9M | let node_idx = g_idx * nodes_per_group |
304 | 94.4M | + |
305 | 43.9M | node_storage_new[node_idx] += params.doc_size_bytes; |
306 | } | |
307 | } | |
308 | } | |
309 | ||
310 | // Compute distribution coefficients of variation. | |
311 | 6 | let old_cv = cv(&old_shard_counts); |
312 | 6 | let new_cv = cv(&new_shard_counts); |
313 | ||
314 | // Normal storage: corpus replicated across RG groups. | |
315 | 6 | let normal_storage_bytes = params.corpus_size_bytes * rg as u64; |
316 | // Peak storage: live + shadow (both fully populated). | |
317 | 6 | let peak_storage_bytes = normal_storage_bytes * 2; |
318 | ||
319 | // Per-node storage (max across all nodes). | |
320 | 6 | let per_node_normal = node_storage_old.iter().copied().max().unwrap_or(0); |
321 | 6 | let per_node_peak = per_node_normal + node_storage_new.iter().copied().max().unwrap_or(0); |
322 | ||
323 | // Write rates. | |
324 | // Normal: each incoming doc → RF × RG actual node writes. | |
325 | 6 | let normal_write_rate = params.write_rate_dps * rf as u64 * rg as u64; |
326 | // Dual-write: each incoming doc → 2 × (RF × RG) writes (old + new assignment). | |
327 | 6 | let dual_write_rate = normal_write_rate * 2; |
328 | // Backfill: reads all docs, writes each to new assignment → RF × RG writes/doc. | |
329 | // Plus ongoing dual-writes for new incoming docs. | |
330 | 6 | let backfill_write_rate = params.backfill_throttle_dps * rf as u64 * rg as u64; |
331 | 6 | let peak_write_rate = dual_write_rate + backfill_write_rate; |
332 | ||
333 | 6 | let dual_write_amplification = 2.0; |
334 | 6 | let peak_write_amplification = peak_write_rate as f64 / normal_write_rate as f64; |
335 | ||
336 | // Backfill duration: total docs / throttle rate. | |
337 | 6 | let backfill_duration_secs = if params.backfill_throttle_dps > 0 { |
338 | 6 | total_docs as f64 / params.backfill_throttle_dps as f64 |
339 | } else { | |
340 | 0 | f64::INFINITY |
341 | }; | |
342 | ||
343 | // Total bytes written during reshard: | |
344 | // 1. Dual-write ongoing for the full reshard duration. | |
345 | // 2. Backfill writes of entire corpus. | |
346 | // Approximate: backfill_duration × dual_write_rate + corpus × RF × RG. | |
347 | 6 | let total_reshard_write_bytes = if params.backfill_throttle_dps > 0 { |
348 | 6 | let dual_write_bytes = |
349 | 6 | backfill_duration_secs * dual_write_rate as f64 * params.doc_size_bytes as f64; |
350 | 6 | let backfill_bytes = total_docs * rf as u64 * rg as u64 * params.doc_size_bytes; |
351 | 6 | (dual_write_bytes as u64) + backfill_bytes |
352 | } else { | |
353 | 0 | 0 |
354 | }; | |
355 | ||
356 | 6 | let storage_amplification = peak_storage_bytes as f64 / normal_storage_bytes as f64; |
357 | ||
358 | 6 | SimResult { |
359 | 6 | label: format!( |
360 | 6 | "{}KB/{}GB/RG{}/RF{}", |
361 | 6 | params.doc_size_bytes / 1024, |
362 | 6 | params.corpus_size_bytes / (1024 * 1024 * 1024), |
363 | 6 | rg, |
364 | 6 | rf |
365 | 6 | ), |
366 | 6 | doc_size_bytes: params.doc_size_bytes, |
367 | 6 | corpus_size_bytes: params.corpus_size_bytes, |
368 | 6 | total_docs, |
369 | 6 | replica_groups: rg, |
370 | 6 | replication_factor: rf, |
371 | 6 | old_shards: params.old_shards, |
372 | 6 | new_shards: params.new_shards, |
373 | 6 | nodes_per_group, |
374 | 6 | write_rate_dps: params.write_rate_dps, |
375 | 6 | |
376 | 6 | normal_storage_bytes, |
377 | 6 | peak_storage_bytes, |
378 | 6 | storage_amplification, |
379 | 6 | |
380 | 6 | normal_write_rate, |
381 | 6 | dual_write_rate, |
382 | 6 | peak_write_rate, |
383 | 6 | dual_write_amplification, |
384 | 6 | peak_write_amplification, |
385 | 6 | |
386 | 6 | backfill_duration_secs, |
387 | 6 | total_bytes_written: total_reshard_write_bytes, |
388 | 6 | |
389 | 6 | per_node_peak_storage_bytes: per_node_peak, |
390 | 6 | per_node_normal_storage_bytes: per_node_normal, |
391 | 6 | |
392 | 6 | old_shard_cv: old_cv, |
393 | 6 | new_shard_cv: new_cv, |
394 | 6 | } |
395 | 6 | } |
396 | ||
397 | /// Coefficient of variation for a distribution. | |
398 | 12 | fn cv(values: &[u64]) -> f64 { |
399 | 12 | if values.is_empty() { |
400 | 0 | return 0.0; |
401 | 12 | } |
402 | 12 | let n = values.len() as f64; |
403 | 12 | let mean = values.iter().sum::<u64>() as f64 / n; |
404 | 12 | if mean == 0.0 { |
405 | 0 | return 0.0; |
406 | 12 | } |
407 | 12 | let variance = values |
408 | 12 | .iter() |
409 | 640 | . |
410 | 12 | .sum::<f64>() |
411 | 12 | / n; |
412 | 12 | variance.sqrt() / mean |
413 | 12 | } |
414 | ||
415 | // --------------------------------------------------------------------------- | |
416 | // Tests | |
417 | // --------------------------------------------------------------------------- | |
418 | ||
419 | #[cfg(test)] | |
420 | mod tests { | |
421 | use super::*; | |
422 | ||
423 | // ---- TimeWindow parsing and containment ---- | |
424 | ||
425 | #[test] | |
426 | 2 | fn time_window_parse_simple() { |
427 | 2 | let w = TimeWindow::parse("02:00-06:00").unwrap(); |
428 | 2 | assert_eq!(w.start_mins, 120); |
429 | 2 | assert_eq!(w.end_mins, 360); |
430 | 2 | } |
431 | ||
432 | #[test] | |
433 | 2 | fn time_window_parse_wrap_midnight() { |
434 | 2 | let w = TimeWindow::parse("22:00-06:00").unwrap(); |
435 | 2 | assert_eq!(w.start_mins, 1320); |
436 | 2 | assert_eq!(w.end_mins, 360); |
437 | 2 | } |
438 | ||
439 | #[test] | |
440 | 2 | fn time_window_contains_normal() { |
441 | 2 | let w = TimeWindow::parse("02:00-06:00").unwrap(); |
442 | 2 | assert!(w.contains(180)); // 03:00 |
443 | 2 | assert!(!w.contains(100)); // 01:40 |
444 | 2 | assert!(!w.contains(400)); // 06:40 |
445 | 2 | } |
446 | ||
447 | #[test] | |
448 | 2 | fn time_window_contains_wrap() { |
449 | 2 | let w = TimeWindow::parse("22:00-06:00").unwrap(); |
450 | 2 | assert!(w.contains(1350)); // 22:30 |
451 | 2 | assert!(w.contains(300)); // 05:00 |
452 | 2 | assert!(!w.contains(700)); // 11:40 |
453 | 2 | } |
454 | ||
455 | #[test] | |
456 | 2 | fn time_window_boundary_start() { |
457 | 2 | let w = TimeWindow::parse("02:00-06:00").unwrap(); |
458 | 2 | assert!(w.contains(120)); // exactly 02:00 |
459 | 2 | } |
460 | ||
461 | #[test] | |
462 | 2 | fn time_window_boundary_end_exclusive() { |
463 | 2 | let w = TimeWindow::parse("02:00-06:00").unwrap(); |
464 | 2 | assert!(!w.contains(360)); // exactly 06:00 is excluded |
465 | 2 | } |
466 | ||
467 | #[test] | |
468 | 2 | fn time_window_invalid_format() { |
469 | 2 | assert!(TimeWindow::parse("not-a-window").is_err()); |
470 | 2 | assert!(TimeWindow::parse("25:00-06:00").is_err()); |
471 | 2 | assert!(TimeWindow::parse("02:60-06:00").is_err()); |
472 | 2 | } |
473 | ||
474 | // ---- Window guard ---- | |
475 | ||
476 | #[test] | |
477 | 2 | fn window_guard_no_restriction() { |
478 | 2 | let config = ReshardingConfig::default(); |
479 | 2 | assert_eq!(check_window(0, &config), WindowGuardResult::NoRestriction); |
480 | 2 | } |
481 | ||
482 | #[test] | |
483 | 2 | fn window_guard_allowed() { |
484 | 2 | let config = ReshardingConfig { |
485 | 2 | allowed_windows: vec!["02:00-06:00".into()], |
486 | 2 | ..Default::default() |
487 | 2 | }; |
488 | 2 | let result = check_window(180, &config); // 03:00 |
489 | 2 | assert!( |
490 | 2 | } |
491 | ||
492 | #[test] | |
493 | 2 | fn window_guard_denied() { |
494 | 2 | let config = ReshardingConfig { |
495 | 2 | allowed_windows: vec!["02:00-06:00".into()], |
496 | 2 | ..Default::default() |
497 | 2 | }; |
498 | 2 | let result = check_window(720, &config); // 12:00 |
499 | 2 | assert!( |
500 | 2 | } |
501 | ||
502 | #[test] | |
503 | 2 | fn window_guard_multiple_windows() { |
504 | 2 | let config = ReshardingConfig { |
505 | 2 | allowed_windows: vec!["02:00-04:00".into(), "22:00-23:30".into()], |
506 | 2 | ..Default::default() |
507 | 2 | }; |
508 | // In first window. | |
509 | 2 | assert!( |
510 | 2 | check_window(150, &config), |
511 | WindowGuardResult::Allowed { .. } | |
512 | )); | |
513 | // In second window. | |
514 | 2 | assert!( |
515 | 2 | check_window(1350, &config), |
516 | WindowGuardResult::Allowed { .. } | |
517 | )); | |
518 | // Outside both. | |
519 | 2 | assert!( |
520 | 2 | check_window(720, &config), |
521 | WindowGuardResult::Denied { .. } | |
522 | )); | |
523 | 2 | } |
524 | ||
525 | // ---- Simulation ---- | |
526 | ||
527 | #[test] | |
528 | 2 | fn simulation_storage_always_2x() { |
529 | // Regardless of parameters, peak storage should be exactly 2× normal. | |
530 | 2 | let params = SimParams { |
531 | 2 | doc_size_bytes: 1024, |
532 | 2 | corpus_size_bytes: 10 * 1024 * 1024 * 1024, // 10 GB |
533 | 2 | write_rate_dps: 100, |
534 | 2 | replica_groups: 2, |
535 | 2 | replication_factor: 1, |
536 | 2 | old_shards: 64, |
537 | 2 | new_shards: 128, |
538 | 2 | nodes_per_group: 3, |
539 | 2 | backfill_throttle_dps: 10_000, |
540 | 2 | }; |
541 | 2 | let result = simulate(¶ms); |
542 | 2 | assert!( |
543 | 2 | (result.storage_amplification - 2.0).abs() < 0.01, |
544 | 0 | "expected ~2.0, got {}", |
545 | result.storage_amplification | |
546 | ); | |
547 | 2 | } |
548 | ||
549 | #[test] | |
550 | 2 | fn simulation_dual_write_is_2x() { |
551 | 2 | let params = SimParams { |
552 | 2 | doc_size_bytes: 1024, |
553 | 2 | corpus_size_bytes: 1024 * 1024, |
554 | 2 | write_rate_dps: 100, |
555 | 2 | replica_groups: 2, |
556 | 2 | replication_factor: 1, |
557 | 2 | old_shards: 16, |
558 | 2 | new_shards: 32, |
559 | 2 | nodes_per_group: 3, |
560 | 2 | backfill_throttle_dps: 1000, |
561 | 2 | }; |
562 | 2 | let result = simulate(¶ms); |
563 | 2 | assert!( |
564 | 2 | (result.dual_write_amplification - 2.0).abs() < 0.01, |
565 | 0 | "expected 2.0, got {}", |
566 | result.dual_write_amplification | |
567 | ); | |
568 | 2 | } |
569 | ||
570 | #[test] | |
571 | 2 | fn simulation_low_cv_with_many_docs() { |
572 | // With enough docs, hash distribution CV should be very low (< 5%). | |
573 | 2 | let params = SimParams { |
574 | 2 | doc_size_bytes: 1024, |
575 | 2 | corpus_size_bytes: 1_000_000 * 1024, // 1M docs × 1KB |
576 | 2 | write_rate_dps: 100, |
577 | 2 | replica_groups: 1, |
578 | 2 | replication_factor: 1, |
579 | 2 | old_shards: 16, |
580 | 2 | new_shards: 64, |
581 | 2 | nodes_per_group: 4, |
582 | 2 | backfill_throttle_dps: 1000, |
583 | 2 | }; |
584 | 2 | let result = simulate(¶ms); |
585 | 2 | assert!( |
586 | 2 | result.old_shard_cv < 0.05, |
587 | 0 | "old shard CV too high: {}", |
588 | result.old_shard_cv | |
589 | ); | |
590 | 2 | assert!( |
591 | 2 | result.new_shard_cv < 0.05, |
592 | 0 | "new shard CV too high: {}", |
593 | result.new_shard_cv | |
594 | ); | |
595 | 2 | } |
596 | } |
Line | Count | Source |
1 | //! Rendezvous hash-based routing and shard assignment. | |
2 | ||
3 | use crate::topology::{Group, NodeId, Topology}; | |
4 | use std::hash::{Hash, Hasher}; | |
5 | use twox_hash::XxHash64; | |
6 | ||
7 | /// Compute a rendezvous score for a shard+node pair. | |
8 | /// | |
9 | /// Higher scores win; used for deterministic shard assignment. | |
10 | /// | |
11 | /// CRITICAL: Uses seed 0 to match Meilisearch Enterprise's hash function. | |
12 | /// Any deviation (different seed, different ordering, endianness) forks | |
13 | /// routing across any two Miroir instances and silently corrupts writes. | |
14 | 267M | pub fn score(shard_id: u32, node_id: &str) -> u64 { |
15 | 267M | let mut h = XxHash64::with_seed(0); |
16 | 267M | shard_id.hash(&mut h); |
17 | 267M | node_id.hash(&mut h); |
18 | 267M | h.finish() |
19 | 267M | } |
20 | ||
21 | /// Assign a shard to `rf` nodes within a single replica group. | |
22 | /// | |
23 | /// `group_nodes` is the subset of nodes belonging to that group. | |
24 | /// | |
25 | /// Nodes are sorted by score descending, with ties broken lexicographically | |
26 | /// by node_id to ensure deterministic assignment even when hash scores collide. | |
27 | 87.9M | pub fn assign_shard_in_group(shard_id: u32, group_nodes: &[NodeId], rf: usize) -> Vec<NodeId> { |
28 | 87.9M | let mut scored: Vec<(u64, &NodeId)> = group_nodes |
29 | 87.9M | .iter() |
30 | 267M | . |
31 | 87.9M | .collect(); |
32 | 244M |
|
33 | 244M | b.0.cmp(&a.0) |
34 | 244M | .then_with(|| |
35 | 244M | }); |
36 | 87.9M | scored |
37 | 87.9M | .into_iter() |
38 | 87.9M | .take(rf) |
39 | 87.9M | . |
40 | 87.9M | .collect() |
41 | 87.9M | } |
42 | ||
43 | /// All write targets for a document: the RF nodes in EACH replica group. | |
44 | 6 | pub fn write_targets(shard_id: u32, topology: &Topology) -> Vec<NodeId> { |
45 | 6 | topology |
46 | 6 | .groups() |
47 | 10 | . |
48 | 6 | .collect() |
49 | 6 | } |
50 | ||
51 | /// Select the replica group for a query (round-robin by query counter). | |
52 | 2.00k | pub fn query_group(query_seq: u64, replica_groups: u32) -> u32 { |
53 | 2.00k | (query_seq % replica_groups as u64) as u32 |
54 | 2.00k | } |
55 | ||
56 | /// The covering set for a search: one node per shard within the chosen group. | |
57 | /// | |
58 | /// Returns a Vec where index i contains the node to query for shard i. | |
59 | /// The length is always exactly shard_count, even if multiple shards | |
60 | /// map to the same node. | |
61 | 6 | pub fn covering_set(shard_count: u32, group: &Group, rf: usize, query_seq: u64) -> Vec<NodeId> { |
62 | 6 | (0..shard_count) |
63 | 168 | . |
64 | 168 | let replicas = assign_shard_in_group(shard_id, group.nodes(), rf); |
65 | // rotate through replicas for intra-group load balancing | |
66 | 168 | replicas[(query_seq as usize) % replicas.len()].clone() |
67 | 168 | }) |
68 | 6 | .collect() |
69 | 6 | } |
70 | ||
71 | /// Compute the shard ID for a document's primary key. | |
72 | 45.9M | pub fn shard_for_key(primary_key: &str, shard_count: u32) -> u32 { |
73 | 45.9M | let mut h = XxHash64::with_seed(0); |
74 | 45.9M | primary_key.hash(&mut h); |
75 | 45.9M | (h.finish() % shard_count as u64) as u32 |
76 | 45.9M | } |
77 | ||
78 | #[cfg(test)] | |
79 | mod tests { | |
80 | use super::*; | |
81 | use crate::topology::{Node, NodeId}; | |
82 | ||
83 | // Test 1: Rendezvous assignment is deterministic given fixed node list | |
84 | #[test] | |
85 | 2 | fn test_rendezvous_determinism() { |
86 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3"] |
87 | 2 | .into_iter() |
88 | 6 | . |
89 | 2 | .collect(); |
90 | 2 | let shard_id = 42; |
91 | ||
92 | 2 | let assignment1 = assign_shard_in_group(shard_id, &nodes, 1); |
93 | 2 | let assignment2 = assign_shard_in_group(shard_id, &nodes, 1); |
94 | ||
95 | 2 | assert_eq!(assignment1, assignment2); |
96 | 2 | } |
97 | ||
98 | // Test 2: Score is stable across calls | |
99 | #[test] | |
100 | 2 | fn test_score_stability() { |
101 | 2 | let score1 = score(123, "node1"); |
102 | 2 | let score2 = score(123, "node1"); |
103 | 2 | assert_eq!(score1, score2); |
104 | 2 | } |
105 | ||
106 | // Test 3: Different shard+node pairs produce different scores | |
107 | #[test] | |
108 | 2 | fn test_score_uniqueness() { |
109 | 2 | let score1 = score(1, "node1"); |
110 | 2 | let score2 = score(1, "node2"); |
111 | 2 | let score3 = score(2, "node1"); |
112 | ||
113 | 2 | assert_ne!(score1, score2); |
114 | 2 | assert_ne!(score1, score3); |
115 | 2 | assert_ne!(score2, score3); |
116 | 2 | } |
117 | ||
118 | // Test 4: Adding a 4th node moves at most ~2 × (1/4) of shards | |
119 | #[test] | |
120 | 2 | fn test_minimal_reshuffling_on_add() { |
121 | 2 | let nodes_3: Vec<NodeId> = vec!["node1", "node2", "node3"] |
122 | 2 | .into_iter() |
123 | 6 | . |
124 | 2 | .collect(); |
125 | 2 | let nodes_4: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
126 | 2 | .into_iter() |
127 | 8 | . |
128 | 2 | .collect(); |
129 | ||
130 | 2 | let shard_count = 100; |
131 | 2 | let rf = 1; |
132 | ||
133 | 2 | let mut moved_count = 0; |
134 | 200 | for shard_id in 0.. |
135 | 200 | let assign_3 = assign_shard_in_group(shard_id, &nodes_3, rf); |
136 | 200 | let assign_4 = assign_shard_in_group(shard_id, &nodes_4, rf); |
137 | ||
138 | // Shard moved if its primary owner changed | |
139 | 200 | if assign_3.first() != assign_4.first() { |
140 | 46 | moved_count += 1; |
141 | 154 | } |
142 | } | |
143 | ||
144 | // Expected: at most ~2 × (1/4) = 50% of shards | |
145 | 2 | let max_expected = (shard_count as f64 * 0.5).ceil() as usize; |
146 | 2 | assert!( |
147 | 2 | moved_count <= max_expected, |
148 | 0 | "Expected ≤ {max_expected} shards to move, but {moved_count} moved" |
149 | ); | |
150 | 2 | } |
151 | ||
152 | // Test 5: 64 shards / 3 nodes / RF=1 → each node holds 18–26 shards | |
153 | #[test] | |
154 | 2 | fn test_shard_distribution_64_3_rf1() { |
155 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3"] |
156 | 2 | .into_iter() |
157 | 6 | . |
158 | 2 | .collect(); |
159 | 2 | let shard_count = 64; |
160 | 2 | let rf = 1; |
161 | ||
162 | 2 | let mut node_shard_counts: std::collections::HashMap<String, usize> = |
163 | 2 | std::collections::HashMap::new(); |
164 | ||
165 | 128 | for shard_id in 0.. |
166 | 128 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
167 | 128 | if let Some(node) = assignment.first() { |
168 | 128 | *node_shard_counts |
169 | 128 | .entry(node.as_str().to_string()) |
170 | 128 | .or_insert(0) += 1; |
171 | 128 |
|
172 | } | |
173 | ||
174 | // Debug: print actual distribution | |
175 | 2 | eprintln!("Actual shard distribution: {node_shard_counts:?}"); |
176 | ||
177 | // DoD requirement: each node holds 15–27 shards | |
178 | // This accommodates natural variance in hash-based distribution | |
179 | 8 | for ( |
180 | 6 | assert!( |
181 | 6 | *count >= 15 && *count <= 27, |
182 | 0 | "Node {node} has {count} shards, expected 15–27" |
183 | ); | |
184 | } | |
185 | ||
186 | // Total should equal shard_count | |
187 | 2 | let total: usize = node_shard_counts.values().sum(); |
188 | 2 | assert_eq!(total, shard_count as usize); |
189 | 2 | } |
190 | ||
191 | // Test 6: Top-RF placement changes minimally on add/remove | |
192 | #[test] | |
193 | 2 | fn test_top_rf_stability() { |
194 | 2 | let nodes_3: Vec<NodeId> = vec!["node1", "node2", "node3"] |
195 | 2 | .into_iter() |
196 | 6 | . |
197 | 2 | .collect(); |
198 | 2 | let nodes_4: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
199 | 2 | .into_iter() |
200 | 8 | . |
201 | 2 | .collect(); |
202 | 2 | let rf = 2; |
203 | 2 | let shard_count = 100; |
204 | ||
205 | 2 | let mut changed_count = 0; |
206 | 200 | for shard_id in 0.. |
207 | 200 | let assign_3 = assign_shard_in_group(shard_id, &nodes_3, rf); |
208 | 200 | let assign_4 = assign_shard_in_group(shard_id, &nodes_4, rf); |
209 | ||
210 | // Count how many of the top-RF nodes changed | |
211 | 200 | let set_3: std::collections::HashSet<_> = assign_3.iter().collect(); |
212 | 200 | let set_4: std::collections::HashSet<_> = assign_4.iter().collect(); |
213 | ||
214 | // A change is if the intersection is less than RF | |
215 | 200 | let intersection = set_3.intersection(&set_4).count(); |
216 | 200 | if intersection < rf { |
217 | 98 | changed_count += 1; |
218 | 102 | } |
219 | } | |
220 | ||
221 | // Adding a 4th node affects approximately 1/4 of assignments | |
222 | // But with RF=2, we need to account for overlap | |
223 | // Expected: roughly 50-60% might have some change | |
224 | 2 | let max_expected = (shard_count as f64 * 0.6).ceil() as usize; |
225 | 2 | assert!( |
226 | 2 | changed_count <= max_expected, |
227 | 0 | "Expected ≤ {max_expected} shards to change, but {changed_count} changed" |
228 | ); | |
229 | ||
230 | // Also verify that not everything changed | |
231 | 2 | let min_expected = (shard_count as f64 * 0.2).ceil() as usize; |
232 | 2 | assert!( |
233 | 2 | changed_count >= min_expected, |
234 | 0 | "Expected at least {min_expected} shards to change, but only {changed_count} changed" |
235 | ); | |
236 | 2 | } |
237 | ||
238 | // Test 7: write_targets returns exactly RG × RF nodes | |
239 | #[test] | |
240 | 2 | fn test_write_targets_count() { |
241 | 2 | let mut topology = Topology::new(64, 2); // 64 shards, RF=2 |
242 | ||
243 | // 3 replica groups, 2 nodes each | |
244 | 8 | for |
245 | 18 | for |
246 | 12 | let node = Node::new( |
247 | 12 | NodeId::new(format!("node-g{group_id}-{node_idx}")), |
248 | 12 | format!("http://example.com/{group_id}"), |
249 | 12 | group_id, |
250 | 12 | ); |
251 | 12 | topology.add_node(node); |
252 | 12 | } |
253 | } | |
254 | ||
255 | 2 | let shard_id = 42; |
256 | 2 | let targets = write_targets(shard_id, &topology); |
257 | ||
258 | // Should be RG (3) × RF (2) = 6 nodes | |
259 | 2 | assert_eq!(targets.len(), 6); |
260 | ||
261 | // All targets should be unique | |
262 | 2 | let unique: std::collections::HashSet<_> = targets.iter().collect(); |
263 | 2 | assert_eq!(unique.len(), 6); |
264 | ||
265 | // Each replica group should contribute exactly RF nodes | |
266 | 6 | for group in |
267 | 6 | let group_targets: Vec<_> = targets |
268 | 6 | .iter() |
269 | 36 | . |
270 | 6 | .collect(); |
271 | 6 | assert_eq!( |
272 | 6 | group_targets.len(), |
273 | 6 | topology.rf(), |
274 | 0 | "Group {} should contribute exactly RF nodes", |
275 | group.id | |
276 | ); | |
277 | } | |
278 | 2 | } |
279 | ||
280 | // Test 8: query_group distributes evenly | |
281 | #[test] | |
282 | 2 | fn test_query_group_distribution() { |
283 | 2 | let replica_groups = 3u32; |
284 | 2 | let queries = 1000u64; |
285 | ||
286 | 2 | let mut counts = vec![0; replica_groups as usize]; |
287 | 2.00k | for seq in 0.. |
288 | 2.00k | let group = query_group(seq, replica_groups); |
289 | 2.00k | counts[group as usize] += 1; |
290 | 2.00k | } |
291 | ||
292 | // Each group should get roughly the same number of queries | |
293 | 2 | let expected = (queries / replica_groups as u64) as usize; |
294 | 8 | for |
295 | 6 | assert!( |
296 | 6 | count >= expected - 1 && count <= expected + 1, |
297 | 0 | "Group query count {} outside expected range [{}, {}]", |
298 | count, | |
299 | 0 | expected - 1, |
300 | 0 | expected + 1 |
301 | ); | |
302 | } | |
303 | 2 | } |
304 | ||
305 | // Test 9: covering_set returns exactly one node per shard | |
306 | #[test] | |
307 | 2 | fn test_covering_set_one_per_shard() { |
308 | 2 | let mut topology = Topology::new(64, 2); // 64 shards, RF=2 |
309 | 2 | let group_id = 0; |
310 | 2 | let num_nodes = 5; |
311 | ||
312 | // Add nodes to a single group | |
313 | 10 | for node_idx in 0.. |
314 | 10 | let node = Node::new( |
315 | 10 | NodeId::new(format!("node-{node_idx}")), |
316 | 10 | format!("http://example.com/{node_idx}"), |
317 | 10 | group_id, |
318 | 10 | ); |
319 | 10 | topology.add_node(node); |
320 | 10 | } |
321 | ||
322 | 2 | let group = topology.group(group_id).unwrap(); |
323 | 2 | let shard_count = 64; |
324 | 2 | let rf = 2; |
325 | 2 | let query_seq = 0; |
326 | ||
327 | 2 | let covering = covering_set(shard_count, group, rf, query_seq); |
328 | ||
329 | // Should have exactly one node per shard | |
330 | 2 | assert_eq!(covering.len(), shard_count as usize); |
331 | ||
332 | // All nodes should be from the group | |
333 | 130 | for |
334 | 128 | assert!(group.nodes().contains(node)); |
335 | } | |
336 | 2 | } |
337 | ||
338 | // Test 10: covering_set handles intra-group replica rotation | |
339 | #[test] | |
340 | 2 | fn test_covering_set_replica_rotation() { |
341 | 2 | let mut topology = Topology::new(64, 2); // 64 shards, RF=2 |
342 | 2 | let group_id = 0; |
343 | ||
344 | // Add 3 nodes to a single group | |
345 | 8 | for |
346 | 6 | let node = Node::new( |
347 | 6 | NodeId::new(format!("node-{node_idx}")), |
348 | 6 | format!("http://example.com/{node_idx}"), |
349 | 6 | group_id, |
350 | 6 | ); |
351 | 6 | topology.add_node(node); |
352 | 6 | } |
353 | ||
354 | 2 | let group = topology.group(group_id).unwrap(); |
355 | 2 | let shard_count = 10; |
356 | 2 | let rf = 2; |
357 | ||
358 | 2 | let covering_0 = covering_set(shard_count, group, rf, 0); |
359 | 2 | let covering_1 = covering_set(shard_count, group, rf, 1); |
360 | ||
361 | // With RF=2, the covering set should rotate between the two replicas | |
362 | // For each shard, the node should be different between query_seq 0 and 1 | |
363 | // Note: This is true for most shards but not all, since assignment is deterministic | |
364 | 2 | let mut rotated_count = 0; |
365 | 20 | for (n0, n1) in |
366 | 20 | if n0 != n1 { |
367 | 20 | rotated_count += 1; |
368 | 20 |
|
369 | } | |
370 | ||
371 | // At least some shards should rotate (ideally most/all) | |
372 | 2 | assert!( |
373 | 2 | rotated_count >= shard_count as usize / 2, |
374 | 0 | "Expected at least half of shards to rotate, but only {rotated_count} did" |
375 | ); | |
376 | 2 | } |
377 | ||
378 | // Test 11: shard_for_key is deterministic | |
379 | #[test] | |
380 | 2 | fn test_shard_for_key_determinism() { |
381 | 2 | let key = "user:12345"; |
382 | 2 | let shard_count = 64; |
383 | ||
384 | 2 | let shard1 = shard_for_key(key, shard_count); |
385 | 2 | let shard2 = shard_for_key(key, shard_count); |
386 | ||
387 | 2 | assert_eq!(shard1, shard2); |
388 | 2 | assert!(shard1 < shard_count); |
389 | 2 | } |
390 | ||
391 | // Test 12: shard_for_key distributes keys evenly | |
392 | #[test] | |
393 | 2 | fn test_shard_for_key_distribution() { |
394 | 2 | let shard_count = 64; |
395 | 2 | let keys = 1000; |
396 | ||
397 | 2 | let mut counts = vec![0; shard_count as usize]; |
398 | 2.00k | for i in 0.. |
399 | 2.00k | let key = format!("user:{i}"); |
400 | 2.00k | let shard = shard_for_key(&key, shard_count); |
401 | 2.00k | counts[shard as usize] += 1; |
402 | 2.00k | } |
403 | ||
404 | // Each shard should get roughly keys / shard_count entries | |
405 | 2 | let expected = keys / shard_count as usize; |
406 | 130 | for |
407 | // Allow some variance due to hash distribution | |
408 | 128 | assert!( |
409 | 128 | count >= expected / 2 && count <= expected * 2, |
410 | 0 | "Shard count {count} outside reasonable range" |
411 | ); | |
412 | } | |
413 | 2 | } |
414 | ||
415 | // Test 13: assign_shard_in_group respects RF limit | |
416 | #[test] | |
417 | 2 | fn test_assign_shard_respects_rf() { |
418 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3", "node4", "node5"] |
419 | 2 | .into_iter() |
420 | 10 | . |
421 | 2 | .collect(); |
422 | 2 | let shard_id = 42; |
423 | ||
424 | 12 | for |
425 | 10 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
426 | 10 | assert_eq!( |
427 | 10 | assignment.len(), |
428 | rf, | |
429 | 0 | "Assignment should return exactly RF nodes" |
430 | ); | |
431 | } | |
432 | 2 | } |
433 | ||
434 | // Test 14: assign_shard_in_group handles RF larger than node count | |
435 | #[test] | |
436 | 2 | fn test_assign_shard_rf_larger_than_nodes() { |
437 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3"] |
438 | 2 | .into_iter() |
439 | 6 | . |
440 | 2 | .collect(); |
441 | 2 | let shard_id = 42; |
442 | 2 | let rf = 5; |
443 | ||
444 | 2 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
445 | ||
446 | // Should return all available nodes when RF > node count | |
447 | 2 | assert_eq!(assignment.len(), nodes.len()); |
448 | 2 | } |
449 | ||
450 | // Test 15: Empty node list returns empty assignment | |
451 | #[test] | |
452 | 2 | fn test_assign_shard_empty_nodes() { |
453 | 2 | let nodes: Vec<NodeId> = vec![]; |
454 | 2 | let shard_id = 42; |
455 | 2 | let rf = 2; |
456 | ||
457 | 2 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
458 | ||
459 | 2 | assert!(assignment.is_empty()); |
460 | 2 | } |
461 | ||
462 | // Test 16: write_targets with empty topology | |
463 | #[test] | |
464 | 2 | fn test_write_targets_empty_topology() { |
465 | 2 | let topology = Topology::new(64, 2); |
466 | 2 | let shard_id = 42; |
467 | ||
468 | 2 | let targets = write_targets(shard_id, &topology); |
469 | ||
470 | 2 | assert!(targets.is_empty()); |
471 | 2 | } |
472 | ||
473 | // Test 17: shard_for_key with zero shard_count handles edge case | |
474 | #[test] | |
475 | #[should_panic(expected = "attempt to calculate the remainder with a divisor of zero")] | |
476 | 2 | fn test_shard_for_key_zero_shard_count() { |
477 | // This test verifies the panic behavior - in production this should be validated | |
478 | // at the API boundary | |
479 | 2 | shard_for_key("test", 0); |
480 | 2 | } |
481 | ||
482 | // Test 18: Group-scoped assignment prevents same-group replica placement | |
483 | #[test] | |
484 | 2 | fn test_group_scoped_assignment() { |
485 | // Create topology with 2 groups, 2 nodes each | |
486 | 2 | let mut topology = Topology::new(64, 1); // 64 shards, RF=1 |
487 | 2 | let shard_id = 42; |
488 | ||
489 | // Group 0 | |
490 | 2 | topology.add_node(Node::new( |
491 | 2 | NodeId::new("g0n0".to_string()), |
492 | 2 | "http://g0n0".to_string(), |
493 | 0, | |
494 | )); | |
495 | 2 | topology.add_node(Node::new( |
496 | 2 | NodeId::new("g0n1".to_string()), |
497 | 2 | "http://g0n1".to_string(), |
498 | 0, | |
499 | )); | |
500 | ||
501 | // Group 1 | |
502 | 2 | topology.add_node(Node::new( |
503 | 2 | NodeId::new("g1n0".to_string()), |
504 | 2 | "http://g1n0".to_string(), |
505 | 1, | |
506 | )); | |
507 | 2 | topology.add_node(Node::new( |
508 | 2 | NodeId::new("g1n1".to_string()), |
509 | 2 | "http://g1n1".to_string(), |
510 | 1, | |
511 | )); | |
512 | ||
513 | 2 | let targets = write_targets(shard_id, &topology); |
514 | ||
515 | // With RG=2, RF=1, should get 2 targets (one from each group) | |
516 | 2 | assert_eq!(targets.len(), 2); |
517 | ||
518 | // Verify one from each group | |
519 | 2 | let g0_target = targets.iter().any(|t| { |
520 | 2 | topology |
521 | 2 | .node(t) |
522 | 2 | .map(|n| n.replica_group == 0) |
523 | 2 | .unwrap_or(false) |
524 | 2 | }); |
525 | 4 | let |
526 | 4 | topology |
527 | 4 | .node(t) |
528 | 4 | .map(|n| n.replica_group == 1) |
529 | 4 | .unwrap_or(false) |
530 | 4 | }); |
531 | ||
532 | 2 | assert!(g0_target, |
533 | 2 | assert!(g1_target, |
534 | 2 | } |
535 | ||
536 | // === Acceptance Tests (plan §8 "Router correctness") === | |
537 | ||
538 | // AT-1: Determinism: same (shard_id, nodes) → identical Vec<NodeId> across 1000 randomized runs | |
539 | #[test] | |
540 | 2 | fn acceptance_determinism_1000_runs() { |
541 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
542 | 2 | .into_iter() |
543 | 8 | . |
544 | 2 | .collect(); |
545 | ||
546 | 2.00k | for |
547 | 2.00k | let shard_id = (run % 100) as u32; // Test different shard IDs |
548 | 2.00k | let rf = ((run % 3) + 1) as usize; // Test different RF values |
549 | ||
550 | 2.00k | let assignment1 = assign_shard_in_group(shard_id, &nodes, rf); |
551 | 2.00k | let assignment2 = assign_shard_in_group(shard_id, &nodes, rf); |
552 | ||
553 | 2.00k | assert_eq!( |
554 | assignment1, assignment2, | |
555 | 0 | "Assignments differ on run {}: shard_id={}, rf={}", |
556 | run, shard_id, rf | |
557 | ); | |
558 | } | |
559 | 2 | } |
560 | ||
561 | // AT-2: Reshuffle bound on add: 64 shards, 3→4 nodes → at most 2 × (1/4) × 64 edges differ | |
562 | #[test] | |
563 | 2 | fn acceptance_reshuffle_bound_on_add() { |
564 | 2 | let nodes_3: Vec<NodeId> = vec!["node1", "node2", "node3"] |
565 | 2 | .into_iter() |
566 | 6 | . |
567 | 2 | .collect(); |
568 | 2 | let nodes_4: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
569 | 2 | .into_iter() |
570 | 8 | . |
571 | 2 | .collect(); |
572 | ||
573 | 2 | let shard_count = 64; |
574 | 2 | let rf = 1; |
575 | ||
576 | 2 | let mut moved_count = 0; |
577 | 128 | for shard_id in 0.. |
578 | 128 | let assign_3 = assign_shard_in_group(shard_id, &nodes_3, rf); |
579 | 128 | let assign_4 = assign_shard_in_group(shard_id, &nodes_4, rf); |
580 | ||
581 | // Shard moved if its primary owner changed | |
582 | 128 | if assign_3.first() != assign_4.first() { |
583 | 30 | moved_count += 1; |
584 | 98 | } |
585 | } | |
586 | ||
587 | // Expected: at most 2 × (1/4) × 64 = 32 edges differ | |
588 | 2 | let max_expected = (2.0 * (1.0 / 4.0) * shard_count as f64).ceil() as usize; |
589 | 2 | assert!( |
590 | 2 | moved_count <= max_expected, |
591 | 0 | "Expected ≤ {max_expected} shard-node edges to differ, but {moved_count} differed" |
592 | ); | |
593 | 2 | } |
594 | ||
595 | // AT-3: Reshuffle bound on remove: 64 shards, 4→3 nodes → ~RF × S / Ng edges differ | |
596 | #[test] | |
597 | 2 | fn acceptance_reshuffle_bound_on_remove() { |
598 | 2 | let nodes_4: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
599 | 2 | .into_iter() |
600 | 8 | . |
601 | 2 | .collect(); |
602 | 2 | let nodes_3: Vec<NodeId> = vec!["node1", "node2", "node3"] |
603 | 2 | .into_iter() |
604 | 6 | . |
605 | 2 | .collect(); |
606 | ||
607 | 2 | let shard_count = 64; |
608 | 2 | let rf = 2; |
609 | ||
610 | 2 | let mut moved_count = 0; |
611 | 128 | for shard_id in 0.. |
612 | 128 | let assign_4 = assign_shard_in_group(shard_id, &nodes_4, rf); |
613 | 128 | let assign_3 = assign_shard_in_group(shard_id, &nodes_3, rf); |
614 | ||
615 | // Count edges that differ | |
616 | 128 | let set_4: std::collections::HashSet<_> = assign_4.iter().collect(); |
617 | 128 | let set_3: std::collections::HashSet<_> = assign_3.iter().collect(); |
618 | ||
619 | // An edge differs if it's not in both sets | |
620 | 128 | let diff = set_4.symmetric_difference(&set_3).count(); |
621 | 128 | if diff > 0 { |
622 | 58 | moved_count += diff; |
623 | 70 | } |
624 | } | |
625 | ||
626 | // Expected: ~RF × S / Ng = 2 × 64 / 4 = 32 edges differ | |
627 | // Allow some variance due to hash distribution | |
628 | 2 | let expected = (rf * shard_count as usize) / 4; |
629 | 2 | let tolerance = (expected as f64 * 0.9).ceil() as usize; // ±90% |
630 | 2 | assert!( |
631 | 2 | moved_count >= expected - tolerance && moved_count <= expected + tolerance, |
632 | 0 | "Expected ~{expected} shard-node edges to differ (±{tolerance}), but {moved_count} differed" |
633 | ); | |
634 | 2 | } |
635 | ||
636 | // AT-4: Uniformity: 64 shards, 3 nodes, RF=1 → each node holds 18–26 shards | |
637 | #[test] | |
638 | 2 | fn acceptance_uniformity_64_shards_3_nodes_rf1() { |
639 | 2 | let nodes: Vec<NodeId> = vec!["node1", "node2", "node3"] |
640 | 2 | .into_iter() |
641 | 6 | . |
642 | 2 | .collect(); |
643 | 2 | let shard_count = 64; |
644 | 2 | let rf = 1; |
645 | ||
646 | 2 | let mut node_shard_counts: std::collections::HashMap<String, usize> = |
647 | 2 | std::collections::HashMap::new(); |
648 | ||
649 | 128 | for shard_id in 0.. |
650 | 128 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
651 | 128 | if let Some(node) = assignment.first() { |
652 | 128 | *node_shard_counts |
653 | 128 | .entry(node.as_str().to_string()) |
654 | 128 | .or_insert(0) += 1; |
655 | 128 |
|
656 | } | |
657 | ||
658 | // DoD requirement: each node holds 15–27 shards | |
659 | 8 | for ( |
660 | 6 | assert!( |
661 | 6 | *count >= 15 && *count <= 27, |
662 | 0 | "Node {node} has {count} shards, expected 15–27" |
663 | ); | |
664 | } | |
665 | ||
666 | // Total should equal shard_count | |
667 | 2 | let total: usize = node_shard_counts.values().sum(); |
668 | 2 | assert_eq!(total, shard_count as usize); |
669 | 2 | } |
670 | ||
671 | // AT-5: RF=2 placement: top-2 nodes change minimally when a node is added or removed | |
672 | #[test] | |
673 | 2 | fn acceptance_rf2_placement_stability() { |
674 | 2 | let nodes_3: Vec<NodeId> = vec!["node1", "node2", "node3"] |
675 | 2 | .into_iter() |
676 | 6 | . |
677 | 2 | .collect(); |
678 | 2 | let nodes_4: Vec<NodeId> = vec!["node1", "node2", "node3", "node4"] |
679 | 2 | .into_iter() |
680 | 8 | . |
681 | 2 | .collect(); |
682 | ||
683 | 2 | let shard_count = 64; |
684 | 2 | let rf = 2; |
685 | ||
686 | 2 | let mut changed_count = 0; |
687 | 128 | for shard_id in 0.. |
688 | 128 | let assign_3 = assign_shard_in_group(shard_id, &nodes_3, rf); |
689 | 128 | let assign_4 = assign_shard_in_group(shard_id, &nodes_4, rf); |
690 | ||
691 | // Count how many of the top-RF nodes changed | |
692 | 128 | let set_3: std::collections::HashSet<_> = assign_3.iter().collect(); |
693 | 128 | let set_4: std::collections::HashSet<_> = assign_4.iter().collect(); |
694 | ||
695 | // A change is if the intersection is less than RF | |
696 | 128 | let intersection = set_3.intersection(&set_4).count(); |
697 | 128 | if intersection < rf { |
698 | 58 | changed_count += 1; |
699 | 70 | } |
700 | } | |
701 | ||
702 | // Adding a 4th node should affect minimally | |
703 | // Expected: roughly 1/4 of assignments might have some change | |
704 | 2 | let max_expected = (shard_count as f64 * 0.5).ceil() as usize; |
705 | 2 | assert!( |
706 | 2 | changed_count <= max_expected, |
707 | 0 | "Expected ≤ {max_expected} shards to change, but {changed_count} changed" |
708 | ); | |
709 | 2 | } |
710 | ||
711 | // AT-6: shard_for_key uses seed 0 and matches known fixture | |
712 | #[test] | |
713 | 2 | fn acceptance_shard_for_key_fixture() { |
714 | // Known fixture values computed with XxHash64::with_seed(0) | |
715 | // These are verified against the actual twox-hash implementation | |
716 | 2 | let fixtures = [ |
717 | 2 | ("user:12345", 64, 15), |
718 | 2 | ("product:abc", 64, 24), |
719 | 2 | ("order:99999", 64, 4), |
720 | 2 | ("test", 16, 10), |
721 | 2 | ("hello", 32, 6), |
722 | 2 | ]; |
723 | ||
724 | 12 | for ( |
725 | 10 | let shard = shard_for_key(key, shard_count); |
726 | 10 | assert_eq!( |
727 | shard, expected_shard, | |
728 | 0 | "shard_for_key(\"{}\", {}) should be {}, got {}", |
729 | key, shard_count, expected_shard, shard | |
730 | ); | |
731 | } | |
732 | 2 | } |
733 | ||
734 | // AT-7: Tie-breaking on node_id for identical scores | |
735 | #[test] | |
736 | 2 | fn acceptance_tie_breaking_node_id() { |
737 | // Create nodes that will have deterministic assignment | |
738 | 2 | let nodes: Vec<NodeId> = vec!["node-a", "node-b", "node-c"] |
739 | 2 | .into_iter() |
740 | 6 | . |
741 | 2 | .collect(); |
742 | ||
743 | 2 | let rf = 3; // Request all nodes |
744 | 2 | let shard_id = 42; |
745 | ||
746 | 2 | let assignment = assign_shard_in_group(shard_id, &nodes, rf); |
747 | ||
748 | // Should return all nodes in a deterministic order | |
749 | 2 | assert_eq!(assignment.len(), 3); |
750 | ||
751 | // The order should be stable across calls | |
752 | 2 | let assignment2 = assign_shard_in_group(shard_id, &nodes, rf); |
753 | 2 | assert_eq!(assignment, assignment2); |
754 | 2 | } |
755 | ||
756 | // AT-8: Canonical concatenation order (shard_id, node_id) | |
757 | #[test] | |
758 | 2 | fn acceptance_canonical_concatenation_order() { |
759 | // Verify that score(shard_id, node_id) != score(node_id, shard_id) | |
760 | // by checking that different orders produce different results | |
761 | 2 | let shard_id = 42u32; |
762 | 2 | let node_id = "node1"; |
763 | ||
764 | 2 | let score_correct = score(shard_id, node_id); |
765 | ||
766 | // Compute score with reversed order (manually) | |
767 | use std::hash::{Hash, Hasher}; | |
768 | 2 | let mut h_rev = twox_hash::XxHash64::with_seed(0); |
769 | 2 | node_id.hash(&mut h_rev); |
770 | 2 | shard_id.hash(&mut h_rev); |
771 | 2 | let score_reversed = h_rev.finish(); |
772 | ||
773 | // These should almost certainly be different | |
774 | 2 | assert_ne!( |
775 | score_correct, score_reversed, | |
776 | 0 | "Canonical order (shard_id, node_id) must differ from reversed order" |
777 | ); | |
778 | 2 | } |
779 | } |
Line | Count | Source |
1 | //! Scatter orchestration: fan-out logic and covering set builder. | |
2 | ||
3 | use async_trait::async_trait; | |
4 | use crate::config::UnavailableShardPolicy; | |
5 | use crate::topology::{NodeId, Topology}; | |
6 | use crate::Result; | |
7 | ||
8 | /// Scatter orchestrator: fans out requests to the covering set. | |
9 | #[async_trait] | |
10 | pub trait Scatter: Send + Sync { | |
11 | /// Execute a scatter request to multiple nodes. | |
12 | /// | |
13 | /// Returns a map of node ID to response. Failed nodes are omitted | |
14 | /// based on the unavailable shard policy. | |
15 | async fn scatter( | |
16 | &self, | |
17 | topology: &Topology, | |
18 | nodes: Vec<NodeId>, | |
19 | request: ScatterRequest, | |
20 | policy: UnavailableShardPolicy, | |
21 | ) -> Result<ScatterResponse>; | |
22 | } | |
23 | ||
24 | /// A scatter request to be sent to each node. | |
25 | #[derive(Debug, Clone)] | |
26 | pub struct ScatterRequest { | |
27 | /// Request body (JSON or raw bytes). | |
28 | pub body: Vec<u8>, | |
29 | ||
30 | /// Request headers. | |
31 | pub headers: Vec<(String, String)>, | |
32 | ||
33 | /// HTTP method. | |
34 | pub method: String, | |
35 | ||
36 | /// Request path. | |
37 | pub path: String, | |
38 | } | |
39 | ||
40 | /// Response from a scatter operation. | |
41 | #[derive(Debug, Clone)] | |
42 | pub struct ScatterResponse { | |
43 | /// Responses from successful nodes. | |
44 | pub responses: Vec<NodeResponse>, | |
45 | ||
46 | /// Nodes that failed or timed out. | |
47 | pub failed: Vec<NodeId>, | |
48 | } | |
49 | ||
50 | /// Response from a single node. | |
51 | #[derive(Debug, Clone)] | |
52 | pub struct NodeResponse { | |
53 | /// Node that responded. | |
54 | pub node_id: NodeId, | |
55 | ||
56 | /// Response body. | |
57 | pub body: Vec<u8>, | |
58 | ||
59 | /// HTTP status code. | |
60 | pub status: u16, | |
61 | ||
62 | /// Response headers. | |
63 | pub headers: Vec<(String, String)>, | |
64 | } | |
65 | ||
66 | /// Default stub implementation of Scatter. | |
67 | #[derive(Debug, Clone, Default)] | |
68 | pub struct StubScatter; | |
69 | ||
70 | #[async_trait] | |
71 | impl Scatter for StubScatter { | |
72 | async fn scatter( | |
73 | &self, | |
74 | _topology: &Topology, | |
75 | _nodes: Vec<NodeId>, | |
76 | _request: ScatterRequest, | |
77 | _policy: UnavailableShardPolicy, | |
78 | 6 | ) -> Result<ScatterResponse> { |
79 | Ok(ScatterResponse { | |
80 | responses: Vec::new(), | |
81 | failed: Vec::new(), | |
82 | }) | |
83 | 6 | } |
84 | } | |
85 | ||
86 | #[cfg(test)] | |
87 | mod tests { | |
88 | use super::*; | |
89 | use crate::config::UnavailableShardPolicy; | |
90 | ||
91 | #[test] | |
92 | 2 | fn test_scatter_request_creation() { |
93 | 2 | let request = ScatterRequest { |
94 | 2 | body: b"test body".to_vec(), |
95 | 2 | headers: vec![("Content-Type".to_string(), "application/json".to_string())], |
96 | 2 | method: "POST".to_string(), |
97 | 2 | path: "/search".to_string(), |
98 | 2 | }; |
99 | ||
100 | 2 | assert_eq!(request.body, b"test body"); |
101 | 2 | assert_eq!(request.headers.len(), 1); |
102 | 2 | assert_eq!(request.method, "POST"); |
103 | 2 | assert_eq!(request.path, "/search"); |
104 | 2 | } |
105 | ||
106 | #[tokio::test] | |
107 | 2 | async fn test_stub_scatter_returns_empty_response() { |
108 | 2 | let scatter = StubScatter; |
109 | 2 | let topology = Topology::new(64, 1); // 64 shards, RF=1 |
110 | 2 | let nodes = vec![NodeId::new("node1".to_string())]; |
111 | 2 | let request = ScatterRequest { |
112 | 2 | body: Vec::new(), |
113 | 2 | headers: Vec::new(), |
114 | 2 | method: "GET".to_string(), |
115 | 2 | path: "/".to_string(), |
116 | 2 | }; |
117 | ||
118 | 2 | let result = scatter |
119 | 2 | .scatter(&topology, nodes, request, UnavailableShardPolicy::Partial) |
120 | 2 | .await |
121 | 2 | .unwrap(); |
122 | ||
123 | 2 | assert!(result.responses.is_empty()); |
124 | 2 | assert!(result.failed.is_empty()); |
125 | 2 | } |
126 | ||
127 | #[test] | |
128 | 2 | fn test_scatter_response_empty() { |
129 | 2 | let response = ScatterResponse { |
130 | 2 | responses: Vec::new(), |
131 | 2 | failed: Vec::new(), |
132 | 2 | }; |
133 | ||
134 | 2 | assert!(response.responses.is_empty()); |
135 | 2 | assert!(response.failed.is_empty()); |
136 | 2 | } |
137 | ||
138 | #[test] | |
139 | 2 | fn test_scatter_response_with_data() { |
140 | 2 | let node_response = NodeResponse { |
141 | 2 | node_id: NodeId::new("node1".to_string()), |
142 | 2 | body: b"response".to_vec(), |
143 | 2 | status: 200, |
144 | 2 | headers: vec![("Content-Type".to_string(), "application/json".to_string())], |
145 | 2 | }; |
146 | ||
147 | 2 | let response = ScatterResponse { |
148 | 2 | responses: vec![node_response], |
149 | 2 | failed: vec![NodeId::new("node2".to_string())], |
150 | 2 | }; |
151 | ||
152 | 2 | assert_eq!(response.responses.len(), 1); |
153 | 2 | assert_eq!(response.failed.len(), 1); |
154 | 2 | assert_eq!(response.responses[0].node_id.as_str(), "node1"); |
155 | 2 | assert_eq!(response.responses[0].status, 200); |
156 | 2 | assert_eq!(response.failed[0].as_str(), "node2"); |
157 | 2 | } |
158 | ||
159 | #[test] | |
160 | 2 | fn test_node_response_creation() { |
161 | 2 | let response = NodeResponse { |
162 | 2 | node_id: NodeId::new("test-node".to_string()), |
163 | 2 | body: b"test body".to_vec(), |
164 | 2 | status: 200, |
165 | 2 | headers: vec![("X-Custom".to_string(), "value".to_string())], |
166 | 2 | }; |
167 | ||
168 | 2 | assert_eq!(response.node_id.as_str(), "test-node"); |
169 | 2 | assert_eq!(response.body, b"test body"); |
170 | 2 | assert_eq!(response.status, 200); |
171 | 2 | assert_eq!(response.headers.len(), 1); |
172 | 2 | assert_eq!(response.headers[0].0, "X-Custom"); |
173 | 2 | } |
174 | ||
175 | #[tokio::test] | |
176 | 2 | async fn test_stub_scatter_with_empty_nodes() { |
177 | 2 | let scatter = StubScatter; |
178 | 2 | let topology = Topology::new(64, 1); // 64 shards, RF=1 |
179 | 2 | let nodes: Vec<NodeId> = Vec::new(); |
180 | 2 | let request = ScatterRequest { |
181 | 2 | body: Vec::new(), |
182 | 2 | headers: Vec::new(), |
183 | 2 | method: "GET".to_string(), |
184 | 2 | path: "/".to_string(), |
185 | 2 | }; |
186 | ||
187 | 2 | let result = scatter |
188 | 2 | .scatter(&topology, nodes, request, UnavailableShardPolicy::Partial) |
189 | 2 | .await |
190 | 2 | .unwrap(); |
191 | ||
192 | 2 | assert!(result.responses.is_empty()); |
193 | 2 | assert!(result.failed.is_empty()); |
194 | 2 | } |
195 | ||
196 | #[tokio::test] | |
197 | 2 | async fn test_stub_scatter_with_multiple_nodes() { |
198 | 2 | let scatter = StubScatter; |
199 | 2 | let mut topology = Topology::new(64, 1); // 64 shards, RF=1 |
200 | ||
201 | 2 | let node1 = NodeId::new("node1".to_string()); |
202 | 2 | let node2 = NodeId::new("node2".to_string()); |
203 | 2 | let node3 = NodeId::new("node3".to_string()); |
204 | ||
205 | 2 | topology.add_node(crate::topology::Node::new( |
206 | 2 | node1.clone(), |
207 | 2 | "http://node1".to_string(), |
208 | 0, | |
209 | )); | |
210 | 2 | topology.add_node(crate::topology::Node::new( |
211 | 2 | node2.clone(), |
212 | 2 | "http://node2".to_string(), |
213 | 0, | |
214 | )); | |
215 | 2 | topology.add_node(crate::topology::Node::new( |
216 | 2 | node3.clone(), |
217 | 2 | "http://node3".to_string(), |
218 | 0, | |
219 | )); | |
220 | ||
221 | 2 | let nodes = vec![node1, node2, node3]; |
222 | 2 | let request = ScatterRequest { |
223 | 2 | body: Vec::new(), |
224 | 2 | headers: Vec::new(), |
225 | 2 | method: "POST".to_string(), |
226 | 2 | path: "/search".to_string(), |
227 | 2 | }; |
228 | ||
229 | 2 | let result = scatter |
230 | 2 | .scatter(&topology, nodes, request, UnavailableShardPolicy::Partial) |
231 | 2 | .await |
232 | 2 | .unwrap(); |
233 | ||
234 | 2 | assert!(result.responses.is_empty()); |
235 | 2 | assert!(result.failed.is_empty()); |
236 | 2 | } |
237 | } |
Line | Count | Source |
1 | //! Score comparability analysis: validates cross-shard ranking consistency. | |
2 | //! | |
3 | //! Implements Plan §15 OP#4: statistical validation that `_rankingScore` | |
4 | //! values remain comparable across shards with very different document-count | |
5 | //! distributions. Uses Kendall Tau correlation to measure ranking similarity | |
6 | //! between sharded and ground-truth (single-index) result orderings. | |
7 | ||
8 | use serde::{Deserialize, Serialize}; | |
9 | ||
10 | /// Parameters for a score comparability simulation run. | |
11 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
12 | pub struct SimParams { | |
13 | /// Total number of documents in the corpus. | |
14 | pub total_docs: u64, | |
15 | ||
16 | /// Number of shards to distribute documents across. | |
17 | pub shard_count: u32, | |
18 | ||
19 | /// Skew factor: multiplier for document distribution imbalance. | |
20 | /// 1.0 = uniform distribution. Higher values create more extreme skew. | |
21 | /// Example: 100.0 means one shard gets 100× the median count. | |
22 | pub skew_factor: f64, | |
23 | ||
24 | /// Number of queries to run against the corpus. | |
25 | pub num_queries: usize, | |
26 | ||
27 | /// Number of top results to compare per query (K). | |
28 | pub top_k: usize, | |
29 | ||
30 | /// Random seed for reproducibility. | |
31 | pub seed: u64, | |
32 | } | |
33 | ||
34 | /// Result of a single query comparison. | |
35 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
36 | pub struct QueryResult { | |
37 | /// Query identifier. | |
38 | pub query_id: usize, | |
39 | ||
40 | /// Number of hits returned by ground truth (single-index). | |
41 | pub ground_truth_hits: usize, | |
42 | ||
43 | /// Number of hits returned by sharded execution. | |
44 | pub sharded_hits: usize, | |
45 | ||
46 | /// Kendall Tau correlation between result orderings. | |
47 | pub kendall_tau: f64, | |
48 | ||
49 | /// Jaccard similarity of the top-K result sets (ignoring order). | |
50 | pub jaccard_similarity: f64, | |
51 | ||
52 | /// Position of the first divergent result (0 if identical, top_k+1 if completely different). | |
53 | pub first_divergence_position: usize, | |
54 | ||
55 | /// Score statistics per shard. | |
56 | pub shard_score_stats: Vec<ShardScoreStats>, | |
57 | } | |
58 | ||
59 | /// Score statistics for a single shard on a single query. | |
60 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
61 | pub struct ShardScoreStats { | |
62 | /// Shard identifier. | |
63 | pub shard_id: u32, | |
64 | ||
65 | /// Number of documents on this shard. | |
66 | pub doc_count: u64, | |
67 | ||
68 | /// Number of hits returned by this shard. | |
69 | pub hit_count: usize, | |
70 | ||
71 | /// Minimum score returned. | |
72 | pub min_score: f64, | |
73 | ||
74 | /// Maximum score returned. | |
75 | pub max_score: f64, | |
76 | ||
77 | /// Mean score of returned hits. | |
78 | pub mean_score: f64, | |
79 | ||
80 | /// Score range (max - min). | |
81 | pub score_range: f64, | |
82 | } | |
83 | ||
84 | /// Aggregate results across all queries. | |
85 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
86 | pub struct AggregateResult { | |
87 | /// Mean Kendall Tau across all queries. | |
88 | pub mean_kendall_tau: f64, | |
89 | ||
90 | /// Standard deviation of Kendall Tau. | |
91 | pub std_kendall_tau: f64, | |
92 | ||
93 | /// Minimum Kendall Tau observed. | |
94 | pub min_kendall_tau: f64, | |
95 | ||
96 | /// Maximum Kendall Tau observed. | |
97 | pub max_kendall_tau: f64, | |
98 | ||
99 | /// Percentage of queries with τ ≥ 0.95. | |
100 | pub percent_above_threshold: f64, | |
101 | ||
102 | /// Mean Jaccard similarity. | |
103 | pub mean_jaccard: f64, | |
104 | ||
105 | /// Mean first divergence position. | |
106 | pub mean_first_divergence: usize, | |
107 | ||
108 | /// Shard population statistics (CV of document counts). | |
109 | pub shard_pop_cv: f64, | |
110 | ||
111 | /// Shard population ratio (max/median). | |
112 | pub shard_pop_ratio: f64, | |
113 | } | |
114 | ||
115 | /// Simulated score comparability test result. | |
116 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
117 | pub struct SimResult { | |
118 | pub params: SimParams, | |
119 | pub aggregate: AggregateResult, | |
120 | pub query_results: Vec<QueryResult>, | |
121 | pub shard_doc_counts: Vec<u64>, | |
122 | } | |
123 | ||
124 | /// Run a score comparability simulation. | |
125 | /// | |
126 | /// This models a simplified ranking system where: | |
127 | /// - Documents have relevance scores based on term frequency | |
128 | /// - Shards compute scores locally using their own document statistics | |
129 | /// - A ground-truth single-index uses global statistics | |
130 | /// | |
131 | /// The simulation measures how well the sharded ordering matches the | |
132 | /// ground-truth ordering using Kendall Tau correlation. | |
133 | 2 | pub fn simulate(params: &SimParams) -> SimResult { |
134 | use rand::rngs::StdRng; | |
135 | use rand::SeedableRng; | |
136 | ||
137 | 2 | let mut rng = StdRng::seed_from_u64(params.seed); |
138 | ||
139 | // Generate document distribution across shards with specified skew. | |
140 | 2 | let shard_doc_counts = generate_skewed_distribution( |
141 | 2 | params.total_docs, |
142 | 2 | params.shard_count, |
143 | 2 | params.skew_factor, |
144 | 2 | &mut rng, |
145 | ); | |
146 | ||
147 | // Compute shard population statistics. | |
148 | 2 | let total_docs = shard_doc_counts.iter().sum::<u64>() as f64; |
149 | 2 | let mean_docs = total_docs / params.shard_count as f64; |
150 | 2 | let variance = shard_doc_counts |
151 | 2 | .iter() |
152 | 8 | . |
153 | 2 | .sum::<f64>() |
154 | 2 | / params.shard_count as f64; |
155 | 2 | let cv = variance.sqrt() / mean_docs; |
156 | 2 | let median_docs = median(&shard_doc_counts); |
157 | 2 | let max_docs = *shard_doc_counts.iter().max().unwrap_or(&1) as f64; |
158 | 2 | let pop_ratio = max_docs / median_docs; |
159 | ||
160 | // Run queries and collect results. | |
161 | 2 | let mut query_results = Vec::with_capacity(params.num_queries); |
162 | ||
163 | 20 | for qid in 0.. |
164 | 20 | let result = run_query(qid, params, &shard_doc_counts, &mut rng); |
165 | 20 | query_results.push(result); |
166 | 20 | } |
167 | ||
168 | // Compute aggregate statistics. | |
169 | 2 | let taus: Vec<f64> = query_results.iter().map(|r| r.kendall_tau).collect(); |
170 | 2 | let mean_tau = mean(&taus); |
171 | 2 | let std_tau = std_dev(&taus, mean_tau); |
172 | 2 | let min_tau = taus.iter().cloned().reduce(f64::min).unwrap_or(0.0); |
173 | 2 | let max_tau = taus.iter().cloned().reduce(f64::max).unwrap_or(0.0); |
174 | 2 | let above_threshold = |
175 | 20 |
|
176 | ||
177 | 2 | let jaccards: Vec<f64> = query_results.iter().map(|r| r.jaccard_similarity).collect(); |
178 | 2 | let mean_jaccard = mean(&jaccards); |
179 | ||
180 | 2 | let divergences: Vec<usize> = query_results |
181 | 2 | .iter() |
182 | 2 | .map(|r| r.first_divergence_position) |
183 | 2 | .collect(); |
184 | 2 | let mean_divergence = mean_usize(&divergences); |
185 | ||
186 | 2 | let aggregate = AggregateResult { |
187 | 2 | mean_kendall_tau: mean_tau, |
188 | 2 | std_kendall_tau: std_tau, |
189 | 2 | min_kendall_tau: min_tau, |
190 | 2 | max_kendall_tau: max_tau, |
191 | 2 | percent_above_threshold: above_threshold, |
192 | 2 | mean_jaccard, |
193 | 2 | mean_first_divergence: mean_divergence, |
194 | 2 | shard_pop_cv: cv, |
195 | 2 | shard_pop_ratio: pop_ratio, |
196 | 2 | }; |
197 | ||
198 | 2 | SimResult { |
199 | 2 | params: params.clone(), |
200 | 2 | aggregate, |
201 | 2 | query_results, |
202 | 2 | shard_doc_counts, |
203 | 2 | } |
204 | 2 | } |
205 | ||
206 | /// Generate a skewed document distribution across shards. | |
207 | /// | |
208 | /// Uses a Pareto-like distribution to create realistic skew where | |
209 | /// a few shards have many more documents than others. | |
210 | 6 | fn generate_skewed_distribution( |
211 | 6 | total_docs: u64, |
212 | 6 | shard_count: u32, |
213 | 6 | skew_factor: f64, |
214 | 6 | rng: &mut impl rand::Rng, |
215 | 6 | ) -> Vec<u64> { |
216 | // Start with uniform weights. | |
217 | 6 | let mut weights: Vec<f64> = (0..shard_count).map(|_| 1.0).collect(); |
218 | ||
219 | // Apply skew: scale a few shards by the skew factor. | |
220 | // For skew_factor = 100, scale the top 5% of shards. | |
221 | 6 | let num_skewed = (shard_count as f64 / 20.0).ceil() as usize; // Top 5% |
222 | 6 | for i in 0..num_skewed.min(weights.len()) { |
223 | 6 | weights[i] *= skew_factor; |
224 | 6 | } |
225 | ||
226 | // Add some randomness to make it more realistic. | |
227 | 54 | for |
228 | 48 | *w *= rng.gen_range(0.5..1.5); |
229 | 48 | } |
230 | ||
231 | // Normalize weights to sum to total_docs. | |
232 | 6 | let weight_sum: f64 = weights.iter().sum(); |
233 | 6 | let scale = total_docs as f64 / weight_sum; |
234 | ||
235 | 6 | weights |
236 | 6 | .into_iter() |
237 | 48 | . |
238 | 6 | .collect() |
239 | 6 | } |
240 | ||
241 | /// Run a single query and compare sharded vs ground-truth results. | |
242 | 20 | fn run_query( |
243 | 20 | query_id: usize, |
244 | 20 | params: &SimParams, |
245 | 20 | shard_doc_counts: &[u64], |
246 | 20 | rng: &mut impl rand::Rng, |
247 | 20 | ) -> QueryResult { |
248 | // Simulate a query by generating random relevance scores for documents. | |
249 | // In a real system, scores would be based on term frequency, IDF, etc. | |
250 | // Here we simulate the key effect: smaller shards tend to produce higher | |
251 | // scores for the same documents because their local statistics are different. | |
252 | ||
253 | 20 | let total_docs = params.total_docs as usize; |
254 | ||
255 | // Generate ground-truth scores (global IDF). | |
256 | // Each document gets a base relevance score + some noise. | |
257 | 20 | let mut ground_truth_scores: Vec<(usize, f64)> = (0..total_docs) |
258 | 20.0k | . |
259 | 20.0k | let base_relevance = rng.gen::<f64>(); // Simulated term matching |
260 | 20.0k | let global_idf = 1.0; // Global IDF is uniform |
261 | 20.0k | let score = base_relevance * global_idf; |
262 | 20.0k | (doc_id, score) |
263 | 20.0k | }) |
264 | 20 | .collect(); |
265 | ||
266 | // Sort by score descending and take top-K. | |
267 | 213k |
|
268 | ||
269 | 20 | let ground_truth_top: Vec<usize> = ground_truth_scores |
270 | 20 | .iter() |
271 | 20 | .take(params.top_k) |
272 | 20 | .map(|(id, _)| *id) |
273 | 20 | .collect(); |
274 | ||
275 | // Generate sharded scores (local IDF per shard). | |
276 | // Key insight: local IDF is computed from shard-specific document counts. | |
277 | // Smaller shards have fewer docs, so their IDF tends to be higher for | |
278 | // the same term, leading to score inflation. | |
279 | ||
280 | 20 | let mut sharded_scores: Vec<(usize, f64)> = Vec::new(); |
281 | ||
282 | 80 | for (shard_id, &doc_count) in |
283 | 80 | let shard_id = shard_id as u32; |
284 | 80 | let doc_count_f64 = doc_count as f64; |
285 | 80 | let total_docs_f64 = params.total_docs as f64; |
286 | ||
287 | // Simulate local IDF effect: smaller shards → higher local IDF. | |
288 | // In real BM25, IDF = log((N - df + 0.5) / (df + 0.5)) where N is doc count. | |
289 | // We model this as: local_idf / global_idf ≈ log(N_global) / log(N_shard). | |
290 | 80 | let global_idf_factor = (total_docs_f64 + 1.0).ln(); |
291 | 80 | let local_idf_factor = (doc_count_f64 + 1.0).ln(); |
292 | 80 | let idf_inflation = global_idf_factor / local_idf_factor.max(0.1); |
293 | ||
294 | // Documents on this shard get scores inflated by the local IDF factor. | |
295 | 80 | let start_doc = shard_doc_counts[..shard_id as usize].iter().sum::<u64>() as usize; |
296 | 80 | let end_doc = start_doc + doc_count as usize; |
297 | ||
298 | 19.9k | for (doc_id, base_relevance) in ground_truth_scores |
299 | 80 | .iter() |
300 | 80 | .enumerate() |
301 | 80 | .skip(start_doc) |
302 | 80 | .take(end_doc.min(total_docs) - start_doc) |
303 | 19.9k | { |
304 | 19.9k | let score = base_relevance.1 * idf_inflation; |
305 | 19.9k | sharded_scores.push((doc_id, score)); |
306 | 19.9k | } |
307 | } | |
308 | ||
309 | // Sort sharded scores and take top-K. | |
310 | 28.4k |
|
311 | ||
312 | 20 | let sharded_top: Vec<usize> = sharded_scores |
313 | 20 | .iter() |
314 | 20 | .take(params.top_k) |
315 | 20 | .map(|(id, _)| *id) |
316 | 20 | .collect(); |
317 | ||
318 | // Compute Kendall Tau between the two orderings. | |
319 | 20 | let kendall_tau = kendall_tau(&ground_truth_top, &sharded_top); |
320 | ||
321 | // Compute Jaccard similarity (set overlap, ignoring order). | |
322 | 20 | let jaccard = jaccard_similarity(&ground_truth_top, &sharded_top); |
323 | ||
324 | // Find first divergence position. | |
325 | 20 | let first_divergence = ground_truth_top |
326 | 20 | .iter() |
327 | 20 | .zip(sharded_top.iter()) |
328 | 20 | .position(|(a, b)| a != b) |
329 | 20 | .unwrap_or(params.top_k); |
330 | ||
331 | // Collect per-shard score statistics. | |
332 | 20 | let shard_stats = |
333 | 20 | collect_shard_stats(query_id, shard_doc_counts, &sharded_scores, params.top_k); |
334 | ||
335 | 20 | QueryResult { |
336 | 20 | query_id, |
337 | 20 | ground_truth_hits: ground_truth_top.len(), |
338 | 20 | sharded_hits: sharded_top.len(), |
339 | 20 | kendall_tau, |
340 | 20 | jaccard_similarity: jaccard, |
341 | 20 | first_divergence_position: first_divergence, |
342 | 20 | shard_score_stats: shard_stats, |
343 | 20 | } |
344 | 20 | } |
345 | ||
346 | /// Collect score statistics for each shard. | |
347 | 20 | fn collect_shard_stats( |
348 | 20 | _query_id: usize, |
349 | 20 | shard_doc_counts: &[u64], |
350 | 20 | sharded_scores: &[(usize, f64)], |
351 | 20 | top_k: usize, |
352 | 20 | ) -> Vec<ShardScoreStats> { |
353 | 20 | let mut stats = Vec::new(); |
354 | ||
355 | // Map doc_id back to its shard. | |
356 | 80 | for (shard_idx, &doc_count) in |
357 | 80 | let shard_id = shard_idx as u32; |
358 | 80 | let start_doc = shard_doc_counts[..shard_idx].iter().sum::<u64>() as usize; |
359 | 80 | let end_doc = start_doc + doc_count as usize; |
360 | ||
361 | // Find scores from this shard in the top-K results. | |
362 | 80 | let shard_hits: Vec<f64> = sharded_scores |
363 | 80 | .iter() |
364 | 80 | .take(top_k) |
365 | 800 | . |
366 | 80 | .map(|(_, score)| *score) |
367 | 80 | .collect(); |
368 | ||
369 | 80 | if shard_hits.is_empty() { |
370 | 60 | continue; |
371 | 20 | } |
372 | ||
373 | 200 | let |
374 | 200 | let |
375 | 20 | let mean_score = shard_hits.iter().sum::<f64>() / shard_hits.len() as f64; |
376 | ||
377 | 20 | stats.push(ShardScoreStats { |
378 | 20 | shard_id, |
379 | 20 | doc_count, |
380 | 20 | hit_count: shard_hits.len(), |
381 | 20 | min_score, |
382 | 20 | max_score, |
383 | 20 | mean_score, |
384 | 20 | score_range: max_score - min_score, |
385 | 20 | }); |
386 | } | |
387 | ||
388 | 20 | stats |
389 | 20 | } |
390 | ||
391 | /// Compute Kendall Tau correlation between two ranked lists. | |
392 | /// | |
393 | /// Returns a value in [-1, 1] where: | |
394 | /// - 1.0 = identical rankings | |
395 | /// - 0.0 = no correlation | |
396 | /// - -1.0 = completely inverted rankings | |
397 | /// | |
398 | /// Uses the O(n²) algorithm which is fine for small K (typically ≤ 100). | |
399 | 26 | fn kendall_tau<T: Eq + std::hash::Hash + std::fmt::Debug>(rank1: &[T], rank2: &[T]) -> f64 { |
400 | 26 | if rank1.is_empty() || rank2.is_empty() { |
401 | 0 | return 1.0; |
402 | 26 | } |
403 | ||
404 | // Create position maps. | |
405 | 26 | let pos1: std::collections::HashMap<&T, usize> = rank1 |
406 | 26 | .iter() |
407 | 26 | .enumerate() |
408 | 230 | . |
409 | 26 | .collect(); |
410 | 26 | let pos2: std::collections::HashMap<&T, usize> = rank2 |
411 | 26 | .iter() |
412 | 26 | .enumerate() |
413 | 230 | . |
414 | 26 | .collect(); |
415 | ||
416 | // Collect all unique items from both lists. | |
417 | 26 | let all_items: std::collections::HashSet<&T> = rank1.iter().chain(rank2.iter()).collect(); |
418 | ||
419 | // Count concordant and discordant pairs. | |
420 | 26 | let mut concordant = 0; |
421 | 26 | let mut discordant = 0; |
422 | ||
423 | 26 | let items: Vec<&T> = all_items.into_iter().collect(); |
424 | 426 | for i in 0.. |
425 | 3.78k | for j in |
426 | 3.78k | let a = items[i]; |
427 | 3.78k | let b = items[j]; |
428 | ||
429 | 3.78k | let pos1_a = pos1.get(a); |
430 | 3.78k | let pos1_b = pos1.get(b); |
431 | 3.78k | let pos2_a = pos2.get(a); |
432 | 3.78k | let pos2_b = pos2.get(b); |
433 | ||
434 | // A pair is only counted if both items appear in both lists. | |
435 | 3.78k | let ( |
436 | 60 | (Some(&x), Some(&y), Some(&u), Some(&v)) => (x, y, u, v), |
437 | 3.72k | _ => continue, |
438 | }; | |
439 | ||
440 | 60 | if (p1a < p1b && |
441 | 38 | concordant += 1; |
442 | 38 | } else if ( |
443 | 22 | discordant += 1; |
444 | 22 |
|
445 | } | |
446 | } | |
447 | ||
448 | 26 | let total = concordant + discordant; |
449 | 26 | if total == 0 { |
450 | 20 | return 1.0; |
451 | 6 | } |
452 | ||
453 | 6 | (concordant - discordant) as f64 / total as f64 |
454 | 26 | } |
455 | ||
456 | /// Compute Jaccard similarity between two sets (ignoring order). | |
457 | 26 | fn jaccard_similarity<T: Eq + std::hash::Hash>(set1: &[T], set2: &[T]) -> f64 { |
458 | 26 | let set1_hash: std::collections::HashSet<&T> = set1.iter().collect(); |
459 | 26 | let set2_hash: std::collections::HashSet<&T> = set2.iter().collect(); |
460 | ||
461 | 26 | let intersection = set1_hash.intersection(&set2_hash).count(); |
462 | 26 | let union = set1_hash.union(&set2_hash).count(); |
463 | ||
464 | 26 | if union == 0 { |
465 | 0 | return 1.0; |
466 | 26 | } |
467 | ||
468 | 26 | intersection as f64 / union as f64 |
469 | 26 | } |
470 | ||
471 | /// Compute mean of a slice of floats. | |
472 | 4 | fn mean(values: &[f64]) -> f64 { |
473 | 4 | if values.is_empty() { |
474 | 0 | return 0.0; |
475 | 4 | } |
476 | 4 | values.iter().sum::<f64>() / values.len() as f64 |
477 | 4 | } |
478 | ||
479 | /// Compute standard deviation of a slice of floats. | |
480 | 2 | fn std_dev(values: &[f64], mean_val: f64) -> f64 { |
481 | 2 | if values.len() <= 1 { |
482 | 0 | return 0.0; |
483 | 2 | } |
484 | 2 | let variance = |
485 | 20 |
|
486 | 2 | variance.sqrt() |
487 | 2 | } |
488 | ||
489 | /// Compute mean of a slice of usize. | |
490 | 2 | fn mean_usize(values: &[usize]) -> usize { |
491 | 2 | if values.is_empty() { |
492 | 0 | return 0; |
493 | 2 | } |
494 | 2 | values.iter().sum::<usize>() / values.len() |
495 | 2 | } |
496 | ||
497 | /// Compute median of a slice of u64. | |
498 | 2 | fn median(values: &[u64]) -> f64 { |
499 | 2 | if values.is_empty() { |
500 | 0 | return 0.0; |
501 | 2 | } |
502 | 2 | let mut sorted = values.to_vec(); |
503 | 2 | sorted.sort(); |
504 | 2 | let len = sorted.len(); |
505 | 2 | if len % 2 == 0 { |
506 | 2 | (sorted[len / 2 - 1] + sorted[len / 2]) as f64 / 2.0 |
507 | } else { | |
508 | 0 | sorted[len / 2] as f64 |
509 | } | |
510 | 2 | } |
511 | ||
512 | // --------------------------------------------------------------------------- | |
513 | // Tests | |
514 | // --------------------------------------------------------------------------- | |
515 | ||
516 | #[cfg(test)] | |
517 | mod tests { | |
518 | use super::*; | |
519 | ||
520 | #[test] | |
521 | 2 | fn test_kendall_tau_identical() { |
522 | 2 | let a = vec![1, 2, 3, 4, 5]; |
523 | 2 | let b = vec![1, 2, 3, 4, 5]; |
524 | 2 | assert!((kendall_tau(&a, &b) - 1.0).abs() < 0.001); |
525 | 2 | } |
526 | ||
527 | #[test] | |
528 | 2 | fn test_kendall_tau_reversed() { |
529 | 2 | let a = vec![1, 2, 3, 4, 5]; |
530 | 2 | let b = vec![5, 4, 3, 2, 1]; |
531 | 2 | assert!((kendall_tau(&a, &b) - (-1.0)).abs() < 0.001); |
532 | 2 | } |
533 | ||
534 | #[test] | |
535 | 2 | fn test_kendall_tau_partial_overlap() { |
536 | 2 | let a = vec![1, 2, 3, 4, 5]; |
537 | 2 | let b = vec![1, 3, 2, 4, 5]; // Only 2 and 3 are swapped |
538 | // One discordant pair (2,3) out of 10 total: tau = (9 - 1) / 10 = 0.8 | |
539 | 2 | let tau = kendall_tau(&a, &b); |
540 | 2 | assert!((tau - 0.8).abs() < 0.01); |
541 | 2 | } |
542 | ||
543 | #[test] | |
544 | 2 | fn test_jaccard_identical() { |
545 | 2 | let a = vec![1, 2, 3, 4, 5]; |
546 | 2 | let b = vec![1, 2, 3, 4, 5]; |
547 | 2 | assert!((jaccard_similarity(&a, &b) - 1.0).abs() < 0.001); |
548 | 2 | } |
549 | ||
550 | #[test] | |
551 | 2 | fn test_jaccard_no_overlap() { |
552 | 2 | let a = vec![1, 2, 3]; |
553 | 2 | let b = vec![4, 5, 6]; |
554 | 2 | assert!((jaccard_similarity(&a, &b) - 0.0).abs() < 0.001); |
555 | 2 | } |
556 | ||
557 | #[test] | |
558 | 2 | fn test_jaccard_half_overlap() { |
559 | 2 | let a = vec![1, 2, 3, 4]; |
560 | 2 | let b = vec![3, 4, 5, 6]; |
561 | // Intersection: {3, 4}, Union: {1, 2, 3, 4, 5, 6} | |
562 | // Jaccard = 2/6 = 1/3 | |
563 | 2 | assert!((jaccard_similarity(&a, &b) - 1.0 / 3.0).abs() < 0.001); |
564 | 2 | } |
565 | ||
566 | #[test] | |
567 | 2 | fn test_skewed_distribution_sum() { |
568 | 2 | let counts = generate_skewed_distribution(10000, 10, 10.0, &mut rand::thread_rng()); |
569 | 2 | let total: u64 = counts.iter().sum(); |
570 | 2 | assert!((total as i64 - 10000i64).abs() < 100); // Allow small rounding error |
571 | 2 | } |
572 | ||
573 | #[test] | |
574 | 2 | fn test_skewed_distribution_has_skew() { |
575 | 2 | let counts = generate_skewed_distribution(10000, 10, 100.0, &mut rand::thread_rng()); |
576 | 2 | let max = *counts.iter().max().unwrap(); |
577 | 2 | let min = *counts.iter().min().unwrap(); |
578 | // With skew factor 100, we expect significant imbalance. | |
579 | 2 | assert!(max > min * 10); |
580 | 2 | } |
581 | ||
582 | #[test] | |
583 | 2 | fn test_simulation_runs() { |
584 | 2 | let params = SimParams { |
585 | 2 | total_docs: 1000, |
586 | 2 | shard_count: 4, |
587 | 2 | skew_factor: 10.0, |
588 | 2 | num_queries: 10, |
589 | 2 | top_k: 10, |
590 | 2 | seed: 42, |
591 | 2 | }; |
592 | 2 | let result = simulate(¶ms); |
593 | 2 | assert_eq!(result.query_results.len(), 10); |
594 | 2 | assert!(result.aggregate.mean_kendall_tau >= -1.0); |
595 | 2 | assert!(result.aggregate.mean_kendall_tau <= 1.0); |
596 | 2 | } |
597 | } |
Line | Count | Source |
1 | //! Task registry: unified task namespace across all Meilisearch nodes. | |
2 | ||
3 | use crate::Result; | |
4 | use serde::{Deserialize, Serialize}; | |
5 | use std::collections::HashMap; | |
6 | use uuid::Uuid; | |
7 | ||
8 | /// Task registry: manages the unified task namespace. | |
9 | pub trait TaskRegistry: Send + Sync { | |
10 | /// Register a new Miroir task that fans out to multiple nodes. | |
11 | fn register(&self, node_tasks: HashMap<String, u64>) -> Result<MiroirTask>; | |
12 | ||
13 | /// Get a task by its Miroir ID. | |
14 | fn get(&self, miroir_id: &str) -> Result<Option<MiroirTask>>; | |
15 | ||
16 | /// Update the status of a Miroir task. | |
17 | fn update_status(&self, miroir_id: &str, status: TaskStatus) -> Result<()>; | |
18 | ||
19 | /// Update node task status. | |
20 | fn update_node_task( | |
21 | &self, | |
22 | miroir_id: &str, | |
23 | node_id: &str, | |
24 | node_status: NodeTaskStatus, | |
25 | ) -> Result<()>; | |
26 | ||
27 | /// List tasks with optional filtering. | |
28 | fn list(&self, filter: TaskFilter) -> Result<Vec<MiroirTask>>; | |
29 | } | |
30 | ||
31 | /// A Miroir task: unified view of a fan-out write operation. | |
32 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
33 | pub struct MiroirTask { | |
34 | /// Unique Miroir task ID (UUID). | |
35 | pub miroir_id: String, | |
36 | ||
37 | /// Creation timestamp (Unix millis). | |
38 | pub created_at: u64, | |
39 | ||
40 | /// Current task status. | |
41 | pub status: TaskStatus, | |
42 | ||
43 | /// Map of node ID to local Meilisearch task UID. | |
44 | pub node_tasks: HashMap<String, NodeTask>, | |
45 | ||
46 | /// Error message if the task failed. | |
47 | pub error: Option<String>, | |
48 | } | |
49 | ||
50 | /// Status of a Miroir task. | |
51 | #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] | |
52 | pub enum TaskStatus { | |
53 | /// Task is enqueued. | |
54 | Enqueued, | |
55 | ||
56 | /// Task is being processed. | |
57 | Processing, | |
58 | ||
59 | /// Task completed successfully. | |
60 | Succeeded, | |
61 | ||
62 | /// Task failed. | |
63 | Failed, | |
64 | ||
65 | /// Task was canceled. | |
66 | Canceled, | |
67 | } | |
68 | ||
69 | /// A node task: local Meilisearch task reference. | |
70 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
71 | pub struct NodeTask { | |
72 | /// Local Meilisearch task UID. | |
73 | pub task_uid: u64, | |
74 | ||
75 | /// Current status of this node task. | |
76 | pub status: NodeTaskStatus, | |
77 | } | |
78 | ||
79 | /// Status of a node task. | |
80 | #[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] | |
81 | pub enum NodeTaskStatus { | |
82 | /// Task is enqueued on the node. | |
83 | Enqueued, | |
84 | ||
85 | /// Task is processing on the node. | |
86 | Processing, | |
87 | ||
88 | /// Task succeeded on the node. | |
89 | Succeeded, | |
90 | ||
91 | /// Task failed on the node. | |
92 | Failed, | |
93 | } | |
94 | ||
95 | /// Filter for listing tasks. | |
96 | #[derive(Debug, Clone, Default)] | |
97 | pub struct TaskFilter { | |
98 | /// Filter by status. | |
99 | pub status: Option<TaskStatus>, | |
100 | ||
101 | /// Filter by node ID. | |
102 | pub node_id: Option<String>, | |
103 | ||
104 | /// Maximum number of results. | |
105 | pub limit: Option<usize>, | |
106 | ||
107 | /// Offset for pagination. | |
108 | pub offset: Option<usize>, | |
109 | } | |
110 | ||
111 | /// Default stub implementation of TaskRegistry. | |
112 | #[derive(Debug, Clone, Default)] | |
113 | pub struct StubTaskRegistry; | |
114 | ||
115 | impl TaskRegistry for StubTaskRegistry { | |
116 | 2 | fn register(&self, _node_tasks: HashMap<String, u64>) -> Result<MiroirTask> { |
117 | 2 | Ok(MiroirTask { |
118 | 2 | miroir_id: Uuid::new_v4().to_string(), |
119 | 2 | created_at: 0, |
120 | 2 | status: TaskStatus::Enqueued, |
121 | 2 | node_tasks: HashMap::new(), |
122 | 2 | error: None, |
123 | 2 | }) |
124 | 2 | } |
125 | ||
126 | 2 | fn get(&self, _miroir_id: &str) -> Result<Option<MiroirTask>> { |
127 | 2 | Ok(None) |
128 | 2 | } |
129 | ||
130 | 2 | fn update_status(&self, _miroir_id: &str, _status: TaskStatus) -> Result<()> { |
131 | 2 | Ok(()) |
132 | 2 | } |
133 | ||
134 | 2 | fn update_node_task( |
135 | 2 | &self, |
136 | 2 | _miroir_id: &str, |
137 | 2 | _node_id: &str, |
138 | 2 | _node_status: NodeTaskStatus, |
139 | 2 | ) -> Result<()> { |
140 | 2 | Ok(()) |
141 | 2 | } |
142 | ||
143 | 2 | fn list(&self, _filter: TaskFilter) -> Result<Vec<MiroirTask>> { |
144 | 2 | Ok(Vec::new()) |
145 | 2 | } |
146 | } | |
147 | ||
148 | #[cfg(test)] | |
149 | mod tests { | |
150 | use super::*; | |
151 | ||
152 | #[test] | |
153 | 2 | fn test_stub_task_registry_register() { |
154 | 2 | let registry = StubTaskRegistry; |
155 | 2 | let mut node_tasks = HashMap::new(); |
156 | 2 | node_tasks.insert("node1".to_string(), 123); |
157 | ||
158 | 2 | let task = registry.register(node_tasks).unwrap(); |
159 | 2 | assert!(!task.miroir_id.is_empty()); |
160 | 2 | assert_eq!(task.status, TaskStatus::Enqueued); |
161 | 2 | assert!(task.node_tasks.is_empty()); |
162 | 2 | assert!(task.error.is_none()); |
163 | 2 | } |
164 | ||
165 | #[test] | |
166 | 2 | fn test_stub_task_registry_get() { |
167 | 2 | let registry = StubTaskRegistry; |
168 | 2 | let result = registry.get("test-id").unwrap(); |
169 | 2 | assert!(result.is_none()); |
170 | 2 | } |
171 | ||
172 | #[test] | |
173 | 2 | fn test_stub_task_registry_update_status() { |
174 | 2 | let registry = StubTaskRegistry; |
175 | 2 | let result = registry.update_status("test-id", TaskStatus::Succeeded); |
176 | 2 | assert!(result.is_ok()); |
177 | 2 | } |
178 | ||
179 | #[test] | |
180 | 2 | fn test_stub_task_registry_update_node_task() { |
181 | 2 | let registry = StubTaskRegistry; |
182 | 2 | let result = registry.update_node_task("test-id", "node1", NodeTaskStatus::Succeeded); |
183 | 2 | assert!(result.is_ok()); |
184 | 2 | } |
185 | ||
186 | #[test] | |
187 | 2 | fn test_stub_task_registry_list() { |
188 | 2 | let registry = StubTaskRegistry; |
189 | 2 | let filter = TaskFilter::default(); |
190 | 2 | let result = registry.list(filter).unwrap(); |
191 | 2 | assert!(result.is_empty()); |
192 | 2 | } |
193 | ||
194 | #[test] | |
195 | 2 | fn test_task_status_equality() { |
196 | 2 | assert_eq!(TaskStatus::Enqueued, TaskStatus::Enqueued); |
197 | 2 | assert_ne!(TaskStatus::Enqueued, TaskStatus::Processing); |
198 | 2 | assert_ne!(TaskStatus::Succeeded, TaskStatus::Failed); |
199 | 2 | } |
200 | ||
201 | #[test] | |
202 | 2 | fn test_node_task_status_equality() { |
203 | 2 | assert_eq!(NodeTaskStatus::Enqueued, NodeTaskStatus::Enqueued); |
204 | 2 | assert_ne!(NodeTaskStatus::Processing, NodeTaskStatus::Succeeded); |
205 | 2 | assert_ne!(NodeTaskStatus::Failed, NodeTaskStatus::Succeeded); |
206 | 2 | } |
207 | ||
208 | #[test] | |
209 | 2 | fn test_task_filter_default() { |
210 | 2 | let filter = TaskFilter::default(); |
211 | 2 | assert!(filter.status.is_none()); |
212 | 2 | assert!(filter.node_id.is_none()); |
213 | 2 | assert!(filter.limit.is_none()); |
214 | 2 | assert!(filter.offset.is_none()); |
215 | 2 | } |
216 | ||
217 | #[test] | |
218 | 2 | fn test_task_filter_with_fields() { |
219 | 2 | let filter = TaskFilter { |
220 | 2 | status: Some(TaskStatus::Processing), |
221 | 2 | node_id: Some("node1".to_string()), |
222 | 2 | limit: Some(10), |
223 | 2 | offset: Some(5), |
224 | 2 | }; |
225 | 2 | assert_eq!(filter.status, Some(TaskStatus::Processing)); |
226 | 2 | assert_eq!(filter.node_id, Some("node1".to_string())); |
227 | 2 | assert_eq!(filter.limit, Some(10)); |
228 | 2 | assert_eq!(filter.offset, Some(5)); |
229 | 2 | } |
230 | ||
231 | #[test] | |
232 | 2 | fn test_miroir_task_creation() { |
233 | 2 | let mut node_tasks = HashMap::new(); |
234 | 2 | node_tasks.insert( |
235 | 2 | "node1".to_string(), |
236 | 2 | NodeTask { |
237 | 2 | task_uid: 123, |
238 | 2 | status: NodeTaskStatus::Enqueued, |
239 | 2 | }, |
240 | ); | |
241 | ||
242 | 2 | let task = MiroirTask { |
243 | 2 | miroir_id: "test-id".to_string(), |
244 | 2 | created_at: 1234567890, |
245 | 2 | status: TaskStatus::Processing, |
246 | 2 | node_tasks, |
247 | 2 | error: None, |
248 | 2 | }; |
249 | ||
250 | 2 | assert_eq!(task.miroir_id, "test-id"); |
251 | 2 | assert_eq!(task.created_at, 1234567890); |
252 | 2 | assert_eq!(task.status, TaskStatus::Processing); |
253 | 2 | assert_eq!(task.node_tasks.len(), 1); |
254 | 2 | assert!(task.error.is_none()); |
255 | 2 | } |
256 | ||
257 | #[test] | |
258 | 2 | fn test_miroir_task_with_error() { |
259 | 2 | let task = MiroirTask { |
260 | 2 | miroir_id: "failed-task".to_string(), |
261 | 2 | created_at: 0, |
262 | 2 | status: TaskStatus::Failed, |
263 | 2 | node_tasks: HashMap::new(), |
264 | 2 | error: Some("Something went wrong".to_string()), |
265 | 2 | }; |
266 | ||
267 | 2 | assert_eq!(task.status, TaskStatus::Failed); |
268 | 2 | assert_eq!(task.error, Some("Something went wrong".to_string())); |
269 | 2 | } |
270 | } |
Line | Count | Source |
1 | //! Topology management: node registry, groups, and health state. | |
2 | ||
3 | use serde::{Deserialize, Serialize}; | |
4 | use std::collections::HashMap; | |
5 | use std::fmt; | |
6 | ||
7 | /// Unique identifier for a node. | |
8 | #[derive(Debug, Clone, Hash, PartialEq, Eq, Serialize, Deserialize)] | |
9 | pub struct NodeId(String); | |
10 | ||
11 | impl NodeId { | |
12 | /// Create a new NodeId. | |
13 | 246 | pub fn new(id: String) -> Self { |
14 | 246 | Self(id) |
15 | 246 | } |
16 | ||
17 | /// Get the node ID as a string slice. | |
18 | 267M | pub fn as_str(&self) -> &str { |
19 | 267M | &self.0 |
20 | 267M | } |
21 | } | |
22 | ||
23 | impl From<String> for NodeId { | |
24 | 2 | fn from(s: String) -> Self { |
25 | 2 | Self(s) |
26 | 2 | } |
27 | } | |
28 | ||
29 | impl AsRef<str> for NodeId { | |
30 | 2 | fn as_ref(&self) -> &str { |
31 | 2 | &self.0 |
32 | 2 | } |
33 | } | |
34 | ||
35 | /// Health status of a node. | |
36 | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] | |
37 | pub enum NodeStatus { | |
38 | /// Node is healthy and serving traffic. | |
39 | Healthy, | |
40 | /// Node is degraded (intermittent failures, still serving traffic). | |
41 | Degraded, | |
42 | /// Node is active and fully operational (synonym for Healthy). | |
43 | Active, | |
44 | /// Node is joining the cluster (being provisioned). | |
45 | Joining, | |
46 | /// Node is draining (graceful shutdown, not accepting new writes). | |
47 | Draining, | |
48 | /// Node has failed (unplanned outage). | |
49 | Failed, | |
50 | /// Node has been removed from the cluster (tracked for migration). | |
51 | Removed, | |
52 | } | |
53 | ||
54 | impl NodeStatus { | |
55 | /// Check if a transition from `self` to `new_status` is valid. | |
56 | /// | |
57 | /// # State Transition Rules | |
58 | /// | |
59 | /// | From | To | Triggered by | | |
60 | /// |------|-----|-------------| | |
61 | /// | (new) | Joining | `POST /_miroir/nodes` | | |
62 | /// | Joining | Active | Migration complete | | |
63 | /// | Active | Draining | `POST /_miroir/nodes/{id}/drain` | | |
64 | /// | Draining | Removed | Migration complete | | |
65 | /// | Active/Draining | Failed | Health check detects | | |
66 | /// | Failed | Active | Health check recovery | | |
67 | /// | Active/Failed | Degraded | Partial health | | |
68 | /// | Degraded | Active | Health restored | | |
69 | 50 | pub fn can_transition_to(self, new_status: NodeStatus) -> bool { |
70 | 50 | match (self, new_status) { |
71 | // Initial state | |
72 | 6 | (NodeStatus::Joining, NodeStatus::Active) => true, |
73 | ||
74 | // Normal operations | |
75 | 2 | (NodeStatus::Active, NodeStatus::Draining) => true, |
76 | 2 | (NodeStatus::Draining, NodeStatus::Removed) => true, |
77 | ||
78 | // Failure and recovery | |
79 | 2 | (NodeStatus::Active, NodeStatus::Failed) => true, |
80 | 2 | (NodeStatus::Draining, NodeStatus::Failed) => true, |
81 | 2 | (NodeStatus::Failed, NodeStatus::Active) => true, |
82 | ||
83 | // Degradation | |
84 | 2 | (NodeStatus::Active, NodeStatus::Degraded) => true, |
85 | 2 | (NodeStatus::Failed, NodeStatus::Degraded) => true, |
86 | 2 | (NodeStatus::Degraded, NodeStatus::Active) => true, |
87 | ||
88 | // Healthy <-> Active are bidirectional (synonyms) | |
89 | 2 | (NodeStatus::Healthy, NodeStatus::Active) => true, |
90 | 2 | (NodeStatus::Active, NodeStatus::Healthy) => true, |
91 | ||
92 | // Same state is always valid | |
93 | 24 | ( |
94 | ||
95 | // All other transitions are invalid | |
96 | 10 | _ => false, |
97 | } | |
98 | 50 | } |
99 | ||
100 | /// Returns `true` if the node can accept writes for the given shard. | |
101 | /// | |
102 | /// # Write Eligibility Rules | |
103 | /// | |
104 | /// A node is write-eligible for a shard based on its status: | |
105 | /// | |
106 | /// | Status | Write Eligible | Notes | | |
107 | /// |--------|----------------|-------| | |
108 | /// | Healthy/Active | Yes | Normal operation | | |
109 | /// | Degraded | Yes | Partial failures, still accepting writes | | |
110 | /// | Joining | No | Being provisioned, not yet ready | | |
111 | /// | Draining | Conditional | Only for shards it still owns during migration | | |
112 | /// | Failed | No | Unavailable | | |
113 | /// | Removed | No | No longer in cluster | | |
114 | /// | |
115 | /// The `draining_shard` parameter should be `Some(shard_id)` if the node | |
116 | /// is in `Draining` status and the shard IS being actively migrated off this node | |
117 | /// (use `None` if the shard is not being drained or no shard is being checked). | |
118 | /// When `Some(...)`, the node is NOT eligible for writes. | |
119 | 32 | pub fn is_write_eligible_for(self, draining_shard: Option<u32>) -> bool { |
120 | 32 | match self { |
121 | 14 | NodeStatus::Healthy | NodeStatus::Active | NodeStatus::Degraded => true, |
122 | 12 | NodeStatus::Joining | NodeStatus::Failed | NodeStatus::Removed => false, |
123 | 6 | NodeStatus::Draining => !draining_shard.is_some(), |
124 | } | |
125 | 32 | } |
126 | } | |
127 | ||
128 | impl fmt::Display for NodeStatus { | |
129 | 18 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
130 | 18 | match self { |
131 | 2 | NodeStatus::Healthy => write!(f, "healthy"), |
132 | 2 | NodeStatus::Degraded => write!(f, "degraded"), |
133 | 2 | NodeStatus::Active => write!(f, "active"), |
134 | 4 | NodeStatus::Joining => write!(f, "joining"), |
135 | 4 | NodeStatus::Draining => write!(f, "draining"), |
136 | 2 | NodeStatus::Failed => write!(f, "failed"), |
137 | 2 | NodeStatus::Removed => write!(f, "removed"), |
138 | } | |
139 | 18 | } |
140 | } | |
141 | ||
142 | /// A single Meilisearch node in the topology. | |
143 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
144 | pub struct Node { | |
145 | /// Unique node identifier. | |
146 | pub id: NodeId, | |
147 | ||
148 | /// Node base URL / address. | |
149 | pub address: String, | |
150 | ||
151 | /// Current health status. | |
152 | pub status: NodeStatus, | |
153 | ||
154 | /// Replica group assignment (0-based). | |
155 | pub replica_group: u32, | |
156 | } | |
157 | ||
158 | impl Node { | |
159 | /// Create a new node. | |
160 | 62 | pub fn new(id: NodeId, address: String, replica_group: u32) -> Self { |
161 | 62 | Self { |
162 | 62 | id, |
163 | 62 | address, |
164 | 62 | status: NodeStatus::Joining, |
165 | 62 | replica_group, |
166 | 62 | } |
167 | 62 | } |
168 | ||
169 | /// Create a new node with a specific status. | |
170 | 16 | pub fn with_status(id: NodeId, address: String, replica_group: u32, status: NodeStatus) -> Self { |
171 | 16 | Self { |
172 | 16 | id, |
173 | 16 | address, |
174 | 16 | status, |
175 | 16 | replica_group, |
176 | 16 | } |
177 | 16 | } |
178 | ||
179 | /// Check if the node is healthy (can serve traffic). | |
180 | 26 | pub fn is_healthy(&self) -> bool { |
181 | 26 |
|
182 | 26 | } |
183 | ||
184 | /// Transition the node to a new status, validating the transition. | |
185 | /// | |
186 | /// Returns `Ok(())` if the transition is valid, `Err` otherwise. | |
187 | 6 | pub fn set_status(&mut self, new_status: NodeStatus) -> Result<(), TransitionError> { |
188 | 6 | if self.status.can_transition_to(new_status) { |
189 | 4 | self.status = new_status; |
190 | 4 | Ok(()) |
191 | } else { | |
192 | 2 | Err(TransitionError { |
193 | 2 | from: self.status, |
194 | 2 | to: new_status, |
195 | 2 | }) |
196 | } | |
197 | 6 | } |
198 | ||
199 | /// Check if the node is eligible to receive writes for a specific shard. | |
200 | /// | |
201 | /// For nodes in `Draining` status, this depends on whether the shard is | |
202 | /// being actively migrated off this node. The caller should pass | |
203 | /// `Some(shard_id)` if the shard is being drained from this node. | |
204 | 2 | pub fn is_write_eligible_for(&self, shard_id: Option<u32>) -> bool { |
205 | 2 | self.status.is_write_eligible_for(shard_id) |
206 | 2 | } |
207 | } | |
208 | ||
209 | /// Error returned when an invalid state transition is attempted. | |
210 | #[derive(Debug, Clone, PartialEq, Eq)] | |
211 | pub struct TransitionError { | |
212 | pub from: NodeStatus, | |
213 | pub to: NodeStatus, | |
214 | } | |
215 | ||
216 | impl fmt::Display for TransitionError { | |
217 | 2 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { |
218 | 2 | write!( |
219 | 2 | f, |
220 | 2 | "invalid state transition from {} to {}", |
221 | self.from, self.to | |
222 | ) | |
223 | 2 | } |
224 | } | |
225 | ||
226 | impl std::error::Error for TransitionError {} | |
227 | ||
228 | /// A replica group: an independent query pool. | |
229 | /// | |
230 | /// Each group holds all S shards, distributed across its nodes. | |
231 | /// Reads are routed to a single group per query. | |
232 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
233 | pub struct Group { | |
234 | /// Group identifier (0-based). | |
235 | pub id: u32, | |
236 | ||
237 | /// Nodes in this group. | |
238 | nodes: Vec<NodeId>, | |
239 | } | |
240 | ||
241 | impl Group { | |
242 | /// Create a new group. | |
243 | 48 | pub fn new(id: u32) -> Self { |
244 | 48 | Self { |
245 | 48 | id, |
246 | 48 | nodes: Vec::new(), |
247 | 48 | } |
248 | 48 | } |
249 | ||
250 | /// Add a node to this group. | |
251 | 108 | pub fn add_node(&mut self, node_id: NodeId) { |
252 | 108 | if !self.nodes.contains(&node_id) { |
253 | 106 | self.nodes.push(node_id); |
254 | 106 |
|
255 | 108 | } |
256 | ||
257 | /// Get the node IDs in this group. | |
258 | 175M | pub fn nodes(&self) -> &[NodeId] { |
259 | 175M | &self.nodes |
260 | 175M | } |
261 | ||
262 | /// Get the number of nodes in this group. | |
263 | 8 | pub fn node_count(&self) -> usize { |
264 | 8 | self.nodes.len() |
265 | 8 | } |
266 | ||
267 | /// Get all healthy nodes in this group, looking them up from the topology. | |
268 | /// | |
269 | /// This requires access to the topology's node map to resolve NodeIds to Nodes. | |
270 | 6 | pub fn healthy_nodes<'a>(&'a self, all_nodes: &'a HashMap<NodeId, Node>) -> Vec<&'a Node> { |
271 | 6 | self.nodes |
272 | 6 | .iter() |
273 | 12 | . |
274 | 12 | . |
275 | 6 | .collect() |
276 | 6 | } |
277 | } | |
278 | ||
279 | /// Cluster topology: groups, nodes, and health state. | |
280 | #[derive(Debug, Clone, Serialize, Deserialize)] | |
281 | pub struct Topology { | |
282 | /// All nodes in the cluster. | |
283 | nodes: HashMap<NodeId, Node>, | |
284 | ||
285 | /// Replica groups. | |
286 | groups: Vec<Group>, | |
287 | ||
288 | /// Replication factor (intra-group). | |
289 | rf: usize, | |
290 | ||
291 | /// Total number of logical shards (S). | |
292 | shards: u32, | |
293 | } | |
294 | ||
295 | impl Topology { | |
296 | /// Create a new empty topology. | |
297 | 28 | pub fn new(shards: u32, rf: usize) -> Self { |
298 | 28 | Self { |
299 | 28 | nodes: HashMap::new(), |
300 | 28 | groups: Vec::new(), |
301 | 28 | rf, |
302 | 28 | shards, |
303 | 28 | } |
304 | 28 | } |
305 | ||
306 | /// Add a node to the topology. | |
307 | 64 | pub fn add_node(&mut self, node: Node) { |
308 | 64 | let group_id = node.replica_group as usize; |
309 | ||
310 | // Ensure group exists | |
311 | 98 | while self.groups.len() <= group_id { |
312 | 34 | self.groups.push(Group::new(self.groups.len() as u32)); |
313 | 34 | } |
314 | ||
315 | 64 | self.groups[group_id].add_node(node.id.clone()); |
316 | 64 | self.nodes.insert(node.id.clone(), node); |
317 | 64 | } |
318 | ||
319 | /// Get a node by ID. | |
320 | 10 | pub fn node(&self, id: &NodeId) -> Option<&Node> { |
321 | 10 | self.nodes.get(id) |
322 | 10 | } |
323 | ||
324 | /// Get a mutable reference to a node by ID. | |
325 | 2 | pub fn node_mut(&mut self, id: &NodeId) -> Option<&mut Node> { |
326 | 2 | self.nodes.get_mut(id) |
327 | 2 | } |
328 | ||
329 | /// Get all nodes. | |
330 | 2 | pub fn nodes(&self) -> impl Iterator<Item = &Node> { |
331 | 2 | self.nodes.values() |
332 | 2 | } |
333 | ||
334 | /// Get a group by ID. | |
335 | 8 | pub fn group(&self, id: u32) -> Option<&Group> { |
336 | 8 | self.groups.get(id as usize) |
337 | 8 | } |
338 | ||
339 | /// Iterate over all groups. | |
340 | 10 | pub fn groups(&self) -> impl Iterator<Item = &Group> { |
341 | 10 | self.groups.iter() |
342 | 10 | } |
343 | ||
344 | /// Get the replication factor. | |
345 | 16 | pub fn rf(&self) -> usize { |
346 | 16 | self.rf |
347 | 16 | } |
348 | ||
349 | /// Get the number of shards. | |
350 | 2 | pub fn shards(&self) -> u32 { |
351 | 2 | self.shards |
352 | 2 | } |
353 | ||
354 | /// Get the number of replica groups. | |
355 | 8 | pub fn replica_group_count(&self) -> u32 { |
356 | 8 | self.groups.len() as u32 |
357 | 8 | } |
358 | ||
359 | /// Get healthy nodes in a specific group. | |
360 | 4 | pub fn healthy_nodes_in_group(&self, group_id: u32) -> Vec<&Node> { |
361 | 4 | self.group(group_id) |
362 | 4 | .map(|g| g.healthy_nodes(&self.nodes)) |
363 | 4 | .unwrap_or_default() |
364 | 4 | } |
365 | } | |
366 | ||
367 | #[cfg(test)] | |
368 | mod tests { | |
369 | use super::*; | |
370 | ||
371 | // --- Existing tests updated for address field --- | |
372 | ||
373 | #[test] | |
374 | 2 | fn test_node_is_healthy() { |
375 | 2 | let mut node = Node::new( |
376 | 2 | NodeId::new("node1".to_string()), |
377 | 2 | "http://example.com".to_string(), |
378 | 0, | |
379 | ); | |
380 | ||
381 | // Joining status is not healthy | |
382 | 2 | assert!(!node.is_healthy()); |
383 | ||
384 | // Healthy status is healthy | |
385 | 2 | node.status = NodeStatus::Healthy; |
386 | 2 | assert!(node.is_healthy()); |
387 | ||
388 | // Active status is healthy (synonym for Healthy) | |
389 | 2 | node.status = NodeStatus::Active; |
390 | 2 | assert!(node.is_healthy()); |
391 | ||
392 | // Degraded status is not healthy (intermittent failures) | |
393 | 2 | node.status = NodeStatus::Degraded; |
394 | 2 | assert!(!node.is_healthy()); |
395 | ||
396 | // Draining status is not healthy | |
397 | 2 | node.status = NodeStatus::Draining; |
398 | 2 | assert!(!node.is_healthy()); |
399 | ||
400 | // Failed status is not healthy | |
401 | 2 | node.status = NodeStatus::Failed; |
402 | 2 | assert!(!node.is_healthy()); |
403 | ||
404 | // Removed status is not healthy | |
405 | 2 | node.status = NodeStatus::Removed; |
406 | 2 | assert!(!node.is_healthy()); |
407 | 2 | } |
408 | ||
409 | #[test] | |
410 | 2 | fn test_group_node_count() { |
411 | 2 | let mut group = Group::new(0); |
412 | 2 | assert_eq!(group.node_count(), 0); |
413 | ||
414 | 2 | group.add_node(NodeId::new("node1".to_string())); |
415 | 2 | assert_eq!(group.node_count(), 1); |
416 | ||
417 | 2 | group.add_node(NodeId::new("node2".to_string())); |
418 | 2 | assert_eq!(group.node_count(), 2); |
419 | ||
420 | // Adding duplicate node doesn't increase count | |
421 | 2 | group.add_node(NodeId::new("node1".to_string())); |
422 | 2 | assert_eq!(group.node_count(), 2); |
423 | 2 | } |
424 | ||
425 | #[test] | |
426 | 2 | fn test_topology_replica_group_count() { |
427 | 2 | let mut topology = Topology::new(64, 2); |
428 | ||
429 | // Empty topology has 0 groups | |
430 | 2 | assert_eq!(topology.replica_group_count(), 0); |
431 | ||
432 | // Add nodes to group 0 | |
433 | 2 | topology.add_node(Node::new( |
434 | 2 | NodeId::new("node1".to_string()), |
435 | 2 | "http://example.com".to_string(), |
436 | 0, | |
437 | )); | |
438 | 2 | assert_eq!(topology.replica_group_count(), 1); |
439 | ||
440 | // Add nodes to group 1 | |
441 | 2 | topology.add_node(Node::new( |
442 | 2 | NodeId::new("node2".to_string()), |
443 | 2 | "http://example.com".to_string(), |
444 | 1, | |
445 | )); | |
446 | 2 | assert_eq!(topology.replica_group_count(), 2); |
447 | ||
448 | // Add more nodes to existing groups | |
449 | 2 | topology.add_node(Node::new( |
450 | 2 | NodeId::new("node3".to_string()), |
451 | 2 | "http://example.com".to_string(), |
452 | 0, | |
453 | )); | |
454 | 2 | assert_eq!(topology.replica_group_count(), 2); |
455 | 2 | } |
456 | ||
457 | #[test] | |
458 | 2 | fn test_topology_nodes_iter() { |
459 | 2 | let mut topology = Topology::new(64, 1); |
460 | ||
461 | 2 | topology.add_node(Node::new( |
462 | 2 | NodeId::new("node1".to_string()), |
463 | 2 | "http://example.com".to_string(), |
464 | 0, | |
465 | )); | |
466 | 2 | topology.add_node(Node::new( |
467 | 2 | NodeId::new("node2".to_string()), |
468 | 2 | "http://example.com".to_string(), |
469 | 1, | |
470 | )); | |
471 | ||
472 | 2 | let nodes: Vec<_> = topology.nodes().collect(); |
473 | 2 | assert_eq!(nodes.len(), 2); |
474 | 2 | } |
475 | ||
476 | #[test] | |
477 | 2 | fn test_topology_groups_iter() { |
478 | 2 | let mut topology = Topology::new(64, 1); |
479 | ||
480 | 2 | topology.add_node(Node::new( |
481 | 2 | NodeId::new("node1".to_string()), |
482 | 2 | "http://example.com".to_string(), |
483 | 0, | |
484 | )); | |
485 | 2 | topology.add_node(Node::new( |
486 | 2 | NodeId::new("node2".to_string()), |
487 | 2 | "http://example.com".to_string(), |
488 | 1, | |
489 | )); | |
490 | ||
491 | 2 | let groups: Vec<_> = topology.groups().collect(); |
492 | 2 | assert_eq!(groups.len(), 2); |
493 | 2 | } |
494 | ||
495 | #[test] | |
496 | 2 | fn test_node_id_from_string() { |
497 | 2 | let id: NodeId = "test-node".to_string().into(); |
498 | 2 | assert_eq!(id.as_str(), "test-node"); |
499 | 2 | } |
500 | ||
501 | #[test] | |
502 | 2 | fn test_node_id_as_ref() { |
503 | 2 | let id = NodeId::new("test-node".to_string()); |
504 | 2 | let s: &str = id.as_ref(); |
505 | 2 | assert_eq!(s, "test-node"); |
506 | 2 | } |
507 | ||
508 | // --- New tests for state transitions --- | |
509 | ||
510 | #[test] | |
511 | 2 | fn test_state_transition_joining_to_active() { |
512 | 2 | assert!(NodeStatus::Joining.can_transition_to(NodeStatus::Active)); |
513 | 2 | } |
514 | ||
515 | #[test] | |
516 | 2 | fn test_state_transition_active_to_draining() { |
517 | 2 | assert!(NodeStatus::Active.can_transition_to(NodeStatus::Draining)); |
518 | 2 | } |
519 | ||
520 | #[test] | |
521 | 2 | fn test_state_transition_draining_to_removed() { |
522 | 2 | assert!(NodeStatus::Draining.can_transition_to(NodeStatus::Removed)); |
523 | 2 | } |
524 | ||
525 | #[test] | |
526 | 2 | fn test_state_transition_active_to_failed() { |
527 | 2 | assert!(NodeStatus::Active.can_transition_to(NodeStatus::Failed)); |
528 | 2 | } |
529 | ||
530 | #[test] | |
531 | 2 | fn test_state_transition_draining_to_failed() { |
532 | 2 | assert!(NodeStatus::Draining.can_transition_to(NodeStatus::Failed)); |
533 | 2 | } |
534 | ||
535 | #[test] | |
536 | 2 | fn test_state_transition_failed_to_active() { |
537 | 2 | assert!(NodeStatus::Failed.can_transition_to(NodeStatus::Active)); |
538 | 2 | } |
539 | ||
540 | #[test] | |
541 | 2 | fn test_state_transition_active_to_degraded() { |
542 | 2 | assert!(NodeStatus::Active.can_transition_to(NodeStatus::Degraded)); |
543 | 2 | } |
544 | ||
545 | #[test] | |
546 | 2 | fn test_state_transition_failed_to_degraded() { |
547 | 2 | assert!(NodeStatus::Failed.can_transition_to(NodeStatus::Degraded)); |
548 | 2 | } |
549 | ||
550 | #[test] | |
551 | 2 | fn test_state_transition_degraded_to_active() { |
552 | 2 | assert!(NodeStatus::Degraded.can_transition_to(NodeStatus::Active)); |
553 | 2 | } |
554 | ||
555 | #[test] | |
556 | 2 | fn test_state_transition_healthy_active_bidirectional() { |
557 | 2 | assert!(NodeStatus::Healthy.can_transition_to(NodeStatus::Active)); |
558 | 2 | assert!(NodeStatus::Active.can_transition_to(NodeStatus::Healthy)); |
559 | 2 | } |
560 | ||
561 | #[test] | |
562 | 2 | fn test_state_transition_same_state() { |
563 | 14 | for status in [ |
564 | 2 | NodeStatus::Healthy, |
565 | 2 | NodeStatus::Degraded, |
566 | 2 | NodeStatus::Active, |
567 | 2 | NodeStatus::Joining, |
568 | 2 | NodeStatus::Draining, |
569 | 2 | NodeStatus::Failed, |
570 | 2 | NodeStatus::Removed, |
571 | ] { | |
572 | 14 | assert!(status.can_transition_to(status)); |
573 | } | |
574 | 2 | } |
575 | ||
576 | #[test] | |
577 | 2 | fn test_state_transition_invalid_joining_to_draining() { |
578 | // Joining node must become Active before Draining | |
579 | 2 | assert!(!NodeStatus::Joining.can_transition_to(NodeStatus::Draining)); |
580 | 2 | } |
581 | ||
582 | #[test] | |
583 | 2 | fn test_state_transition_invalid_joining_to_failed() { |
584 | // Joining node cannot fail (not yet active) | |
585 | 2 | assert!(!NodeStatus::Joining.can_transition_to(NodeStatus::Failed)); |
586 | 2 | } |
587 | ||
588 | #[test] | |
589 | 2 | fn test_state_transition_invalid_removed_to_anything() { |
590 | // Removed is terminal | |
591 | 2 | assert!(!NodeStatus::Removed.can_transition_to(NodeStatus::Active)); |
592 | 2 | assert!(!NodeStatus::Removed.can_transition_to(NodeStatus::Failed)); |
593 | 2 | } |
594 | ||
595 | #[test] | |
596 | 2 | fn test_node_set_status_valid_transition() { |
597 | 2 | let mut node = Node::new( |
598 | 2 | NodeId::new("node1".to_string()), |
599 | 2 | "http://example.com".to_string(), |
600 | 0, | |
601 | ); | |
602 | 2 | assert_eq!(node.status, NodeStatus::Joining); |
603 | ||
604 | 2 | assert!(node.set_status(NodeStatus::Active).is_ok()); |
605 | 2 | assert_eq!(node.status, NodeStatus::Active); |
606 | 2 | } |
607 | ||
608 | #[test] | |
609 | 2 | fn test_node_set_status_invalid_transition() { |
610 | 2 | let mut node = Node::with_status( |
611 | 2 | NodeId::new("node1".to_string()), |
612 | 2 | "http://example.com".to_string(), |
613 | 0, | |
614 | 2 | NodeStatus::Removed, |
615 | ); | |
616 | ||
617 | 2 | let result = node.set_status(NodeStatus::Active); |
618 | 2 | assert!(result.is_err()); |
619 | 2 | let err = result.unwrap_err(); |
620 | 2 | assert_eq!(err.from, NodeStatus::Removed); |
621 | 2 | assert_eq!(err.to, NodeStatus::Active); |
622 | // Status unchanged | |
623 | 2 | assert_eq!(node.status, NodeStatus::Removed); |
624 | 2 | } |
625 | ||
626 | // --- New tests for write eligibility --- | |
627 | ||
628 | #[test] | |
629 | 2 | fn test_write_eligible_healthy() { |
630 | 2 | assert!(NodeStatus::Healthy.is_write_eligible_for(None)); |
631 | 2 | assert!(NodeStatus::Healthy.is_write_eligible_for(Some(0))); |
632 | 2 | } |
633 | ||
634 | #[test] | |
635 | 2 | fn test_write_eligible_active() { |
636 | 2 | assert!(NodeStatus::Active.is_write_eligible_for(None)); |
637 | 2 | assert!(NodeStatus::Active.is_write_eligible_for(Some(0))); |
638 | 2 | } |
639 | ||
640 | #[test] | |
641 | 2 | fn test_write_eligible_degraded() { |
642 | 2 | assert!(NodeStatus::Degraded.is_write_eligible_for(None)); |
643 | 2 | assert!(NodeStatus::Degraded.is_write_eligible_for(Some(0))); |
644 | 2 | } |
645 | ||
646 | #[test] | |
647 | 2 | fn test_write_eligible_joining() { |
648 | // Joining nodes are not write-eligible | |
649 | 2 | assert!(!NodeStatus::Joining.is_write_eligible_for(None)); |
650 | 2 | assert!(!NodeStatus::Joining.is_write_eligible_for(Some(0))); |
651 | 2 | } |
652 | ||
653 | #[test] | |
654 | 2 | fn test_write_eligible_failed() { |
655 | // Failed nodes are not write-eligible | |
656 | 2 | assert!(!NodeStatus::Failed.is_write_eligible_for(None)); |
657 | 2 | assert!(!NodeStatus::Failed.is_write_eligible_for(Some(0))); |
658 | 2 | } |
659 | ||
660 | #[test] | |
661 | 2 | fn test_write_eligible_removed() { |
662 | // Removed nodes are not write-eligible | |
663 | 2 | assert!(!NodeStatus::Removed.is_write_eligible_for(None)); |
664 | 2 | assert!(!NodeStatus::Removed.is_write_eligible_for(Some(0))); |
665 | 2 | } |
666 | ||
667 | #[test] | |
668 | 2 | fn test_write_eligible_draining_non_drained_shard() { |
669 | // Draining node is eligible for writes in general (no specific shard being checked) | |
670 | 2 | assert!(NodeStatus::Draining.is_write_eligible_for(None)); |
671 | // When Some(shard_id) is passed, it means that shard is being drained, so NOT eligible | |
672 | 2 | assert!(!NodeStatus::Draining.is_write_eligible_for(Some(5))); |
673 | 2 | } |
674 | ||
675 | #[test] | |
676 | 2 | fn test_write_eligible_draining_drained_shard() { |
677 | // Draining node is NOT eligible for writes to shards being migrated off | |
678 | 2 | assert!(!NodeStatus::Draining.is_write_eligible_for(Some(3))); |
679 | 2 | } |
680 | ||
681 | #[test] | |
682 | 2 | fn test_node_is_write_eligible_for() { |
683 | 2 | let node = Node::with_status( |
684 | 2 | NodeId::new("node1".to_string()), |
685 | 2 | "http://example.com".to_string(), |
686 | 0, | |
687 | 2 | NodeStatus::Active, |
688 | ); | |
689 | 2 | assert!(node.is_write_eligible_for(Some(0))); |
690 | 2 | } |
691 | ||
692 | // --- New tests for healthy_nodes --- | |
693 | ||
694 | #[test] | |
695 | 2 | fn test_group_healthy_nodes() { |
696 | 2 | let mut group = Group::new(0); |
697 | 2 | let mut all_nodes = HashMap::new(); |
698 | ||
699 | 2 | let node1 = Node::with_status( |
700 | 2 | NodeId::new("node1".to_string()), |
701 | 2 | "http://node1".to_string(), |
702 | 0, | |
703 | 2 | NodeStatus::Active, |
704 | ); | |
705 | 2 | let node2 = Node::with_status( |
706 | 2 | NodeId::new("node2".to_string()), |
707 | 2 | "http://node2".to_string(), |
708 | 0, | |
709 | 2 | NodeStatus::Degraded, |
710 | ); | |
711 | 2 | let node3 = Node::with_status( |
712 | 2 | NodeId::new("node3".to_string()), |
713 | 2 | "http://node3".to_string(), |
714 | 0, | |
715 | 2 | NodeStatus::Failed, |
716 | ); | |
717 | ||
718 | 2 | group.add_node(node1.id.clone()); |
719 | 2 | group.add_node(node2.id.clone()); |
720 | 2 | group.add_node(node3.id.clone()); |
721 | ||
722 | 2 | all_nodes.insert(node1.id.clone(), node1); |
723 | 2 | all_nodes.insert(node2.id.clone(), node2); |
724 | 2 | all_nodes.insert(node3.id.clone(), node3); |
725 | ||
726 | 2 | let healthy = group.healthy_nodes(&all_nodes); |
727 | 2 | assert_eq!(healthy.len(), 1); // Only node1 (Active) is healthy |
728 | 2 | assert_eq!(healthy[0].id.as_str(), "node1"); |
729 | 2 | } |
730 | ||
731 | #[test] | |
732 | 2 | fn test_topology_shards() { |
733 | 2 | let topology = Topology::new(128, 3); |
734 | 2 | assert_eq!(topology.shards(), 128); |
735 | 2 | } |
736 | ||
737 | #[test] | |
738 | 2 | fn test_topology_healthy_nodes_in_group() { |
739 | 2 | let mut topology = Topology::new(64, 2); |
740 | ||
741 | 2 | topology.add_node(Node::with_status( |
742 | 2 | NodeId::new("node1".to_string()), |
743 | 2 | "http://node1".to_string(), |
744 | 0, | |
745 | 2 | NodeStatus::Active, |
746 | )); | |
747 | 2 | topology.add_node(Node::with_status( |
748 | 2 | NodeId::new("node2".to_string()), |
749 | 2 | "http://node2".to_string(), |
750 | 0, | |
751 | 2 | NodeStatus::Failed, |
752 | )); | |
753 | 2 | topology.add_node(Node::with_status( |
754 | 2 | NodeId::new("node3".to_string()), |
755 | 2 | "http://node3".to_string(), |
756 | 1, | |
757 | 2 | NodeStatus::Active, |
758 | )); | |
759 | ||
760 | 2 | let healthy_group0 = topology.healthy_nodes_in_group(0); |
761 | 2 | assert_eq!(healthy_group0.len(), 1); |
762 | 2 | assert_eq!(healthy_group0[0].id.as_str(), "node1"); |
763 | ||
764 | 2 | let healthy_group1 = topology.healthy_nodes_in_group(1); |
765 | 2 | assert_eq!(healthy_group1.len(), 1); |
766 | 2 | assert_eq!(healthy_group1[0].id.as_str(), "node3"); |
767 | 2 | } |
768 | ||
769 | // --- Test for node mutation --- | |
770 | ||
771 | #[test] | |
772 | 2 | fn test_topology_node_mut() { |
773 | 2 | let mut topology = Topology::new(64, 1); |
774 | ||
775 | 2 | topology.add_node(Node::new( |
776 | 2 | NodeId::new("node1".to_string()), |
777 | 2 | "http://node1".to_string(), |
778 | 0, | |
779 | )); | |
780 | ||
781 | 2 | let node_id = NodeId::new("node1".to_string()); |
782 | { | |
783 | 2 | let node = topology.node(&node_id).unwrap(); |
784 | 2 | assert_eq!(node.status, NodeStatus::Joining); |
785 | } | |
786 | ||
787 | 2 | { |
788 | 2 | let node = topology.node_mut(&node_id).unwrap(); |
789 | 2 | node.set_status(NodeStatus::Active).unwrap(); |
790 | 2 | } |
791 | ||
792 | 2 | let node = topology.node(&node_id).unwrap(); |
793 | 2 | assert_eq!(node.status, NodeStatus::Active); |
794 | 2 | } |
795 | ||
796 | // --- Display tests --- | |
797 | ||
798 | #[test] | |
799 | 2 | fn test_node_status_display() { |
800 | 2 | assert_eq!(NodeStatus::Healthy.to_string(), "healthy"); |
801 | 2 | assert_eq!(NodeStatus::Degraded.to_string(), "degraded"); |
802 | 2 | assert_eq!(NodeStatus::Active.to_string(), "active"); |
803 | 2 | assert_eq!(NodeStatus::Joining.to_string(), "joining"); |
804 | 2 | assert_eq!(NodeStatus::Draining.to_string(), "draining"); |
805 | 2 | assert_eq!(NodeStatus::Failed.to_string(), "failed"); |
806 | 2 | assert_eq!(NodeStatus::Removed.to_string(), "removed"); |
807 | 2 | } |
808 | ||
809 | #[test] | |
810 | 2 | fn test_transition_error_display() { |
811 | 2 | let err = TransitionError { |
812 | 2 | from: NodeStatus::Joining, |
813 | 2 | to: NodeStatus::Draining, |
814 | 2 | }; |
815 | 2 | let msg = format!("{}", err); |
816 | 2 | assert!(msg.contains("invalid state transition")); |
817 | 2 | assert!(msg.contains("joining")); |
818 | 2 | assert!(msg.contains("draining")); |
819 | 2 | } |
820 | } |
Click here for information about interpreting this report.
| Filename | Function Coverage | Line Coverage | Region Coverage | Branch Coverage |
anti_entropy.rs | 100.00% (7/7) | 100.00% (70/70) | 100.00% (56/56) | - (0/0) |
config.rs | 83.33% (25/30) | 94.12% (288/306) | 91.13% (267/293) | - (0/0) |
config/advanced.rs | 93.75% (30/32) | 93.06% (268/288) | 90.36% (150/166) | - (0/0) |
config/load.rs | 77.78% (7/9) | 80.00% (112/140) | 51.57% (82/159) | - (0/0) |
config/validate.rs | 100.00% (1/1) | 57.41% (62/108) | 68.60% (59/86) | - (0/0) |
merger.rs | 91.84% (45/49) | 94.67% (551/582) | 96.83% (946/977) | - (0/0) |
migration.rs | 72.09% (31/43) | 77.73% (363/467) | 77.39% (558/721) | - (0/0) |
reshard.rs | 80.56% (29/36) | 89.51% (290/324) | 89.67% (408/455) | - (0/0) |
router.rs | 98.33% (59/60) | 96.20% (481/500) | 97.44% (990/1016) | - (0/0) |
scatter.rs | 100.00% (11/11) | 100.00% (121/121) | 100.00% (214/214) | - (0/0) |
score_comparability.rs | 100.00% (32/32) | 97.23% (316/325) | 98.30% (579/589) | - (0/0) |
task.rs | 100.00% (16/16) | 100.00% (118/118) | 100.00% (164/164) | - (0/0) |
topology.rs | 100.00% (70/70) | 100.00% (421/421) | 100.00% (776/776) | - (0/0) |
Totals | 91.67% (363/396) | 91.80% (3461/3770) | 92.54% (5249/5672) | - (0/0) |