1use crate::raft::OxirsNodeId;
7use std::collections::HashMap;
8use std::time::{Duration, SystemTime};
9use tokio::time::sleep;
10
11#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
13pub enum ReplicationStrategy {
14 Synchronous,
16 Asynchronous,
18 SemiSynchronous { min_replicas: usize },
20 RaftConsensus,
22}
23
24impl Default for ReplicationStrategy {
25 fn default() -> Self {
26 Self::RaftConsensus
27 }
28}
29
30#[derive(Debug, Clone)]
32pub struct ReplicaInfo {
33 pub node_id: OxirsNodeId,
35 pub address: String,
37 pub last_applied_index: u64,
39 pub is_healthy: bool,
41 pub last_contact: SystemTime,
43 pub replication_lag: u64,
45 pub latency: Duration,
47}
48
49impl ReplicaInfo {
50 pub fn new(node_id: OxirsNodeId, address: String) -> Self {
52 Self {
53 node_id,
54 address,
55 last_applied_index: 0,
56 is_healthy: true,
57 last_contact: SystemTime::now(),
58 replication_lag: 0,
59 latency: Duration::from_millis(0),
60 }
61 }
62
63 pub fn is_stale(&self, threshold: Duration) -> bool {
65 self.last_contact.elapsed().unwrap_or(Duration::MAX) > threshold
66 }
67
68 pub fn update_health(&mut self, is_healthy: bool) {
70 self.is_healthy = is_healthy;
71 if is_healthy {
72 self.last_contact = SystemTime::now();
73 }
74 }
75}
76
77#[derive(Debug, Clone, Default)]
79pub struct ReplicationStats {
80 pub total_replicas: usize,
81 pub healthy_replicas: usize,
82 pub average_lag: f64,
83 pub max_lag: u64,
84 pub min_lag: u64,
85 pub average_latency: Duration,
86 pub replication_throughput: f64, }
88
89#[derive(Debug)]
91pub struct ReplicationManager {
92 strategy: ReplicationStrategy,
93 replicas: HashMap<OxirsNodeId, ReplicaInfo>,
94 local_node_id: OxirsNodeId,
95 stats: ReplicationStats,
96}
97
98impl ReplicationManager {
99 pub fn new(strategy: ReplicationStrategy, local_node_id: OxirsNodeId) -> Self {
101 Self {
102 strategy,
103 replicas: HashMap::new(),
104 local_node_id,
105 stats: ReplicationStats::default(),
106 }
107 }
108
109 pub fn with_raft_consensus(local_node_id: OxirsNodeId) -> Self {
111 Self::new(ReplicationStrategy::RaftConsensus, local_node_id)
112 }
113
114 pub fn add_replica(&mut self, node_id: OxirsNodeId, address: String) -> bool {
116 if node_id == self.local_node_id {
117 tracing::warn!("Cannot add local node as replica");
118 return false;
119 }
120
121 let replica_info = ReplicaInfo::new(node_id, address.clone());
122 let is_new = !self.replicas.contains_key(&node_id);
123
124 self.replicas.insert(node_id, replica_info);
125
126 if is_new {
127 tracing::info!("Added replica {} at {}", node_id, address);
128 self.update_stats();
129 }
130
131 is_new
132 }
133
134 pub fn remove_replica(&mut self, node_id: OxirsNodeId) -> bool {
136 if let Some(replica) = self.replicas.remove(&node_id) {
137 tracing::info!("Removed replica {} at {}", node_id, replica.address);
138 self.update_stats();
139 true
140 } else {
141 false
142 }
143 }
144
145 pub fn get_replicas(&self) -> &HashMap<OxirsNodeId, ReplicaInfo> {
147 &self.replicas
148 }
149
150 pub fn get_healthy_replicas(&self) -> Vec<&ReplicaInfo> {
152 self.replicas
153 .values()
154 .filter(|replica| replica.is_healthy)
155 .collect()
156 }
157
158 pub fn get_replica(&self, node_id: OxirsNodeId) -> Option<&ReplicaInfo> {
160 self.replicas.get(&node_id)
161 }
162
163 pub fn update_replica_health(&mut self, node_id: OxirsNodeId, is_healthy: bool) -> bool {
165 if let Some(replica) = self.replicas.get_mut(&node_id) {
166 let was_healthy = replica.is_healthy;
167 replica.update_health(is_healthy);
168
169 if was_healthy != is_healthy {
170 tracing::info!(
171 "Replica {} health changed: {} -> {}",
172 node_id,
173 was_healthy,
174 is_healthy
175 );
176 self.update_stats();
177 }
178
179 true
180 } else {
181 false
182 }
183 }
184
185 pub fn update_replica_lag(
187 &mut self,
188 node_id: OxirsNodeId,
189 applied_index: u64,
190 current_index: u64,
191 ) {
192 if let Some(replica) = self.replicas.get_mut(&node_id) {
193 replica.last_applied_index = applied_index;
194 replica.replication_lag = current_index.saturating_sub(applied_index);
195 self.update_stats();
196 }
197 }
198
199 pub async fn health_check(&mut self, stale_threshold: Duration) {
201 let mut changed = false;
202
203 for replica in self.replicas.values_mut() {
204 let was_healthy = replica.is_healthy;
205
206 if replica.is_stale(stale_threshold) {
207 replica.is_healthy = false;
208 }
209
210 if was_healthy != replica.is_healthy {
211 changed = true;
212 tracing::warn!(
213 "Replica {} marked as unhealthy due to staleness",
214 replica.node_id
215 );
216 }
217 }
218
219 if changed {
220 self.update_stats();
221 }
222 }
223
224 pub fn get_strategy(&self) -> &ReplicationStrategy {
226 &self.strategy
227 }
228
229 pub fn set_strategy(&mut self, strategy: ReplicationStrategy) {
231 if self.strategy != strategy {
232 tracing::info!(
233 "Changing replication strategy from {:?} to {:?}",
234 self.strategy,
235 strategy
236 );
237 self.strategy = strategy;
238 }
239 }
240
241 pub fn get_stats(&self) -> &ReplicationStats {
243 &self.stats
244 }
245
246 pub fn is_replication_healthy(&self) -> bool {
248 let healthy_count = self.get_healthy_replicas().len();
249
250 match &self.strategy {
251 ReplicationStrategy::Synchronous => healthy_count == self.replicas.len(),
252 ReplicationStrategy::Asynchronous => true, ReplicationStrategy::SemiSynchronous { min_replicas } => healthy_count >= *min_replicas,
254 ReplicationStrategy::RaftConsensus => {
255 let total_nodes = self.replicas.len() + 1; let majority = (total_nodes / 2) + 1;
258 healthy_count + 1 >= majority }
260 }
261 }
262
263 pub fn required_replica_count(&self) -> usize {
265 match &self.strategy {
266 ReplicationStrategy::Synchronous => self.replicas.len(),
267 ReplicationStrategy::Asynchronous => 0,
268 ReplicationStrategy::SemiSynchronous { min_replicas } => *min_replicas,
269 ReplicationStrategy::RaftConsensus => {
270 let total_nodes = self.replicas.len() + 1;
271 (total_nodes / 2) + 1 - 1 }
273 }
274 }
275
276 fn update_stats(&mut self) {
278 let healthy_replicas_count = self.replicas.values().filter(|r| r.is_healthy).count();
279 let healthy_lags: Vec<u64> = self
280 .replicas
281 .values()
282 .filter(|r| r.is_healthy)
283 .map(|r| r.replication_lag)
284 .collect();
285 let healthy_latencies: Vec<Duration> = self
286 .replicas
287 .values()
288 .filter(|r| r.is_healthy)
289 .map(|r| r.latency)
290 .collect();
291
292 self.stats.total_replicas = self.replicas.len();
293 self.stats.healthy_replicas = healthy_replicas_count;
294
295 if !healthy_lags.is_empty() {
296 let total_lag: u64 = healthy_lags.iter().sum();
297 self.stats.average_lag = total_lag as f64 / healthy_lags.len() as f64;
298 self.stats.max_lag = healthy_lags.iter().copied().max().unwrap_or(0);
299 self.stats.min_lag = healthy_lags.iter().copied().min().unwrap_or(0);
300
301 let total_latency: Duration = healthy_latencies.iter().sum();
302 self.stats.average_latency = total_latency / healthy_latencies.len() as u32;
303 } else {
304 self.stats.average_lag = 0.0;
305 self.stats.max_lag = 0;
306 self.stats.min_lag = 0;
307 self.stats.average_latency = Duration::from_millis(0);
308 }
309 }
310
311 pub async fn maintenance_tick(&mut self, stale_threshold: Duration) {
317 self.health_check(stale_threshold).await;
318
319 if self.stats.total_replicas > 0 {
320 tracing::debug!(
321 "Replication stats: {}/{} healthy, avg lag: {:.1}, max lag: {}",
322 self.stats.healthy_replicas,
323 self.stats.total_replicas,
324 self.stats.average_lag,
325 self.stats.max_lag
326 );
327 }
328 }
329
330 pub async fn run_maintenance(&mut self) {
332 const HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(30);
333 const STALE_THRESHOLD: Duration = Duration::from_secs(60);
334
335 loop {
336 sleep(HEALTH_CHECK_INTERVAL).await;
337
338 self.health_check(STALE_THRESHOLD).await;
339
340 if self.stats.total_replicas > 0 {
342 tracing::debug!(
343 "Replication stats: {}/{} healthy, avg lag: {:.1}, max lag: {}",
344 self.stats.healthy_replicas,
345 self.stats.total_replicas,
346 self.stats.average_lag,
347 self.stats.max_lag
348 );
349 }
350 }
351 }
352}
353
354#[derive(Debug, thiserror::Error)]
356pub enum ReplicationError {
357 #[error("Insufficient replicas: need {required}, have {available}")]
358 InsufficientReplicas { required: usize, available: usize },
359
360 #[error("Replica {node_id} is unhealthy")]
361 UnhealthyReplica { node_id: OxirsNodeId },
362
363 #[error("Replication timeout after {timeout:?}")]
364 Timeout { timeout: Duration },
365
366 #[error("Network error: {message}")]
367 Network { message: String },
368
369 #[error("Serialization error: {message}")]
370 Serialization { message: String },
371}
372
373#[cfg(test)]
374mod tests {
375 use super::*;
376
377 #[test]
378 fn test_replication_strategy_default() {
379 let strategy = ReplicationStrategy::default();
380 assert_eq!(strategy, ReplicationStrategy::RaftConsensus);
381 }
382
383 #[test]
384 fn test_replica_info_creation() {
385 let replica = ReplicaInfo::new(1, "127.0.0.1:8080".to_string());
386
387 assert_eq!(replica.node_id, 1);
388 assert_eq!(replica.address, "127.0.0.1:8080");
389 assert_eq!(replica.last_applied_index, 0);
390 assert!(replica.is_healthy);
391 assert_eq!(replica.replication_lag, 0);
392 assert_eq!(replica.latency, Duration::from_millis(0));
393 }
394
395 #[test]
396 fn test_replica_info_staleness() {
397 let replica = ReplicaInfo::new(1, "127.0.0.1:8080".to_string());
398
399 assert!(!replica.is_stale(Duration::from_secs(10)));
401
402 std::thread::sleep(Duration::from_micros(1));
404
405 assert!(replica.is_stale(Duration::from_nanos(1)));
407 }
408
409 #[test]
410 fn test_replica_info_update_health() {
411 let mut replica = ReplicaInfo::new(1, "127.0.0.1:8080".to_string());
412
413 replica.update_health(false);
415 assert!(!replica.is_healthy);
416
417 replica.update_health(true);
419 assert!(replica.is_healthy);
420 }
421
422 #[test]
423 fn test_replication_manager_creation() {
424 let manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
425
426 assert_eq!(manager.strategy, ReplicationStrategy::Synchronous);
427 assert_eq!(manager.local_node_id, 1);
428 assert!(manager.replicas.is_empty());
429 assert_eq!(manager.stats.total_replicas, 0);
430 }
431
432 #[test]
433 fn test_replication_manager_with_raft_consensus() {
434 let manager = ReplicationManager::with_raft_consensus(1);
435
436 assert_eq!(manager.strategy, ReplicationStrategy::RaftConsensus);
437 assert_eq!(manager.local_node_id, 1);
438 }
439
440 #[test]
441 fn test_replication_manager_add_replica() {
442 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
443
444 assert!(manager.add_replica(2, "127.0.0.1:8081".to_string()));
446 assert_eq!(manager.replicas.len(), 1);
447 assert!(manager.replicas.contains_key(&2));
448
449 assert!(!manager.add_replica(2, "127.0.0.1:8081".to_string()));
451 assert_eq!(manager.replicas.len(), 1);
452
453 assert!(!manager.add_replica(1, "127.0.0.1:8080".to_string()));
455 assert_eq!(manager.replicas.len(), 1);
456 }
457
458 #[tokio::test]
463 async fn regression_maintenance_tick_runs_on_real_manager() {
464 let mut manager = ReplicationManager::new(ReplicationStrategy::RaftConsensus, 1);
465 manager.add_replica(2, "127.0.0.1:8081".to_string());
466 manager.add_replica(3, "127.0.0.1:8082".to_string());
467
468 manager.maintenance_tick(Duration::from_secs(60)).await;
470
471 assert_eq!(manager.get_stats().total_replicas, 2);
472 }
473
474 #[test]
475 fn test_replication_manager_remove_replica() {
476 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
477
478 manager.add_replica(2, "127.0.0.1:8081".to_string());
479 manager.add_replica(3, "127.0.0.1:8082".to_string());
480 assert_eq!(manager.replicas.len(), 2);
481
482 assert!(manager.remove_replica(2));
484 assert_eq!(manager.replicas.len(), 1);
485 assert!(!manager.replicas.contains_key(&2));
486
487 assert!(!manager.remove_replica(4));
489 assert_eq!(manager.replicas.len(), 1);
490 }
491
492 #[test]
493 fn test_replication_manager_get_healthy_replicas() {
494 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
495
496 manager.add_replica(2, "127.0.0.1:8081".to_string());
497 manager.add_replica(3, "127.0.0.1:8082".to_string());
498
499 manager.update_replica_health(3, false);
501
502 let healthy_replicas = manager.get_healthy_replicas();
503 assert_eq!(healthy_replicas.len(), 1);
504 assert_eq!(healthy_replicas[0].node_id, 2);
505 }
506
507 #[test]
508 fn test_replication_manager_update_replica_health() {
509 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
510
511 manager.add_replica(2, "127.0.0.1:8081".to_string());
512
513 assert!(manager.update_replica_health(2, false));
515 let replica = manager.get_replica(2).unwrap();
516 assert!(!replica.is_healthy);
517
518 assert!(!manager.update_replica_health(3, true));
520 }
521
522 #[test]
523 fn test_replication_manager_update_replica_lag() {
524 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
525
526 manager.add_replica(2, "127.0.0.1:8081".to_string());
527
528 manager.update_replica_lag(2, 50, 100);
530 let replica = manager.get_replica(2).unwrap();
531 assert_eq!(replica.last_applied_index, 50);
532 assert_eq!(replica.replication_lag, 50);
533 }
534
535 #[tokio::test]
536 async fn test_replication_manager_health_check() {
537 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
538
539 manager.add_replica(2, "127.0.0.1:8081".to_string());
540 manager.add_replica(3, "127.0.0.1:8082".to_string());
541
542 assert_eq!(manager.get_healthy_replicas().len(), 2);
544
545 manager.health_check(Duration::from_nanos(1)).await;
547 assert_eq!(manager.get_healthy_replicas().len(), 0);
548 }
549
550 #[test]
551 fn test_replication_manager_strategy_change() {
552 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
553
554 assert_eq!(manager.get_strategy(), &ReplicationStrategy::Synchronous);
555
556 manager.set_strategy(ReplicationStrategy::Asynchronous);
557 assert_eq!(manager.get_strategy(), &ReplicationStrategy::Asynchronous);
558 }
559
560 #[test]
561 fn test_replication_manager_health_status() {
562 let mut manager =
563 ReplicationManager::new(ReplicationStrategy::SemiSynchronous { min_replicas: 2 }, 1);
564
565 manager.add_replica(2, "127.0.0.1:8081".to_string());
567 manager.add_replica(3, "127.0.0.1:8082".to_string());
568 manager.add_replica(4, "127.0.0.1:8083".to_string());
569
570 assert!(manager.is_replication_healthy());
572
573 manager.update_replica_health(4, false);
575 assert!(manager.is_replication_healthy());
576
577 manager.update_replica_health(3, false);
579 assert!(!manager.is_replication_healthy());
580 }
581
582 #[test]
583 fn test_replication_manager_required_replica_count() {
584 let mut manager = ReplicationManager::new(ReplicationStrategy::Synchronous, 1);
585 manager.add_replica(2, "127.0.0.1:8081".to_string());
586 manager.add_replica(3, "127.0.0.1:8082".to_string());
587
588 assert_eq!(manager.required_replica_count(), 2);
590
591 manager.set_strategy(ReplicationStrategy::Asynchronous);
593 assert_eq!(manager.required_replica_count(), 0);
594
595 manager.set_strategy(ReplicationStrategy::SemiSynchronous { min_replicas: 1 });
597 assert_eq!(manager.required_replica_count(), 1);
598
599 manager.set_strategy(ReplicationStrategy::RaftConsensus);
601 assert_eq!(manager.required_replica_count(), 1);
603 }
604
605 #[test]
606 fn test_replication_manager_raft_consensus_health() {
607 let mut manager = ReplicationManager::new(ReplicationStrategy::RaftConsensus, 1);
608
609 assert!(manager.is_replication_healthy());
611
612 manager.add_replica(2, "127.0.0.1:8081".to_string());
614 manager.add_replica(3, "127.0.0.1:8082".to_string());
615
616 assert!(manager.is_replication_healthy());
618
619 manager.update_replica_health(3, false);
621 assert!(manager.is_replication_healthy());
622
623 manager.update_replica_health(2, false);
625 assert!(!manager.is_replication_healthy());
626 }
627
628 #[test]
629 fn test_replication_stats_default() {
630 let stats = ReplicationStats::default();
631 assert_eq!(stats.total_replicas, 0);
632 assert_eq!(stats.healthy_replicas, 0);
633 assert_eq!(stats.average_lag, 0.0);
634 assert_eq!(stats.max_lag, 0);
635 assert_eq!(stats.min_lag, 0);
636 assert_eq!(stats.average_latency, Duration::from_millis(0));
637 assert_eq!(stats.replication_throughput, 0.0);
638 }
639
640 #[test]
641 fn test_replication_error_display() {
642 let err = ReplicationError::InsufficientReplicas {
643 required: 3,
644 available: 1,
645 };
646 assert!(err
647 .to_string()
648 .contains("Insufficient replicas: need 3, have 1"));
649
650 let err = ReplicationError::UnhealthyReplica { node_id: 42 };
651 assert!(err.to_string().contains("Replica 42 is unhealthy"));
652
653 let err = ReplicationError::Timeout {
654 timeout: Duration::from_secs(5),
655 };
656 assert!(err.to_string().contains("Replication timeout after 5s"));
657
658 let err = ReplicationError::Network {
659 message: "connection failed".to_string(),
660 };
661 assert!(err.to_string().contains("Network error: connection failed"));
662
663 let err = ReplicationError::Serialization {
664 message: "json error".to_string(),
665 };
666 assert!(err.to_string().contains("Serialization error: json error"));
667 }
668}