Skip to main content

oxirs_cluster/replication/
manager.rs

1//! # Data Replication
2//!
3//! High-level data replication management for distributed RDF storage.
4//! Works with Raft consensus to ensure consistent replication.
5
6use crate::raft::OxirsNodeId;
7use std::collections::HashMap;
8use std::time::{Duration, SystemTime};
9use tokio::time::sleep;
10
11/// Replication strategy for the cluster
12#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
13pub enum ReplicationStrategy {
14    /// Synchronous replication - wait for all replicas
15    Synchronous,
16    /// Asynchronous replication - fire and forget
17    Asynchronous,
18    /// Semi-synchronous - wait for minimum replicas
19    SemiSynchronous { min_replicas: usize },
20    /// Raft consensus - use Raft for replication
21    RaftConsensus,
22}
23
24impl Default for ReplicationStrategy {
25    fn default() -> Self {
26        Self::RaftConsensus
27    }
28}
29
30/// Information about a replica node
31#[derive(Debug, Clone)]
32pub struct ReplicaInfo {
33    /// Unique node identifier
34    pub node_id: OxirsNodeId,
35    /// Address of the replica
36    pub address: String,
37    /// Last successfully applied log index
38    pub last_applied_index: u64,
39    /// Whether the replica is currently healthy
40    pub is_healthy: bool,
41    /// Last time we successfully communicated with this replica
42    pub last_contact: SystemTime,
43    /// Current replication lag in log entries
44    pub replication_lag: u64,
45    /// Network latency to this replica
46    pub latency: Duration,
47}
48
49impl ReplicaInfo {
50    /// Create a new replica info
51    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    /// Check if replica is stale based on contact time
64    pub fn is_stale(&self, threshold: Duration) -> bool {
65        self.last_contact.elapsed().unwrap_or(Duration::MAX) > threshold
66    }
67
68    /// Update health status and contact time
69    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/// Replication statistics
78#[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, // operations per second
87}
88
89/// Replication manager for distributed RDF data
90#[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    /// Create a new replication manager
100    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    /// Create a new replication manager with Raft consensus (default)
110    pub fn with_raft_consensus(local_node_id: OxirsNodeId) -> Self {
111        Self::new(ReplicationStrategy::RaftConsensus, local_node_id)
112    }
113
114    /// Add a replica to the replication set
115    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    /// Remove a replica from the replication set
135    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    /// Get all replicas
146    pub fn get_replicas(&self) -> &HashMap<OxirsNodeId, ReplicaInfo> {
147        &self.replicas
148    }
149
150    /// Get healthy replicas only
151    pub fn get_healthy_replicas(&self) -> Vec<&ReplicaInfo> {
152        self.replicas
153            .values()
154            .filter(|replica| replica.is_healthy)
155            .collect()
156    }
157
158    /// Get replica by node ID
159    pub fn get_replica(&self, node_id: OxirsNodeId) -> Option<&ReplicaInfo> {
160        self.replicas.get(&node_id)
161    }
162
163    /// Update replica health status
164    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    /// Update replica lag information
186    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    /// Check all replicas and mark stale ones as unhealthy
200    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    /// Get the current replication strategy
225    pub fn get_strategy(&self) -> &ReplicationStrategy {
226        &self.strategy
227    }
228
229    /// Change the replication strategy
230    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    /// Get replication statistics
242    pub fn get_stats(&self) -> &ReplicationStats {
243        &self.stats
244    }
245
246    /// Check if replication requirements are met
247    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, // Always considered healthy
253            ReplicationStrategy::SemiSynchronous { min_replicas } => healthy_count >= *min_replicas,
254            ReplicationStrategy::RaftConsensus => {
255                // For Raft, we need a majority of nodes to be healthy
256                let total_nodes = self.replicas.len() + 1; // +1 for local node
257                let majority = (total_nodes / 2) + 1;
258                healthy_count + 1 >= majority // +1 for local node
259            }
260        }
261    }
262
263    /// Get required replica count for the current strategy
264    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 // -1 because local node is not a replica
272            }
273        }
274    }
275
276    /// Update internal statistics
277    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    /// Perform a single maintenance iteration: health-check replicas against
312    /// `stale_threshold` and log current stats. Separated from
313    /// [`ReplicationManager::run_maintenance`]'s infinite loop so a caller can
314    /// drive the cadence itself and re-check a shutdown flag between ticks
315    /// (see `ClusterNode::start_background_tasks`).
316    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    /// Run periodic maintenance tasks
331    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            // Log stats periodically
341            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/// Replication-related errors
355#[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        // Fresh replica should not be stale
400        assert!(!replica.is_stale(Duration::from_secs(10)));
401
402        // Wait a tiny bit to ensure elapsed time passes the threshold
403        std::thread::sleep(Duration::from_micros(1));
404
405        // Simulate old replica by checking against very short threshold
406        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        // Update to unhealthy
414        replica.update_health(false);
415        assert!(!replica.is_healthy);
416
417        // Update to healthy
418        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        // Add replica
445        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        // Adding same replica again should return false
450        assert!(!manager.add_replica(2, "127.0.0.1:8081".to_string()));
451        assert_eq!(manager.replicas.len(), 1);
452
453        // Cannot add local node as replica
454        assert!(!manager.add_replica(1, "127.0.0.1:8080".to_string()));
455        assert_eq!(manager.replicas.len(), 1);
456    }
457
458    /// Regression: a single maintenance iteration must run against the real
459    /// manager (with its actual replicas) and complete, so the background task
460    /// can drive it per-tick while re-checking a shutdown flag between ticks
461    /// (instead of the old infinite loop on a throwaway empty manager).
462    #[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        // One tick must complete and refresh stats over the real replica set.
469        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        // Remove replica
483        assert!(manager.remove_replica(2));
484        assert_eq!(manager.replicas.len(), 1);
485        assert!(!manager.replicas.contains_key(&2));
486
487        // Removing non-existent replica should return false
488        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        // Mark one replica as unhealthy
500        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        // Update health status
514        assert!(manager.update_replica_health(2, false));
515        let replica = manager.get_replica(2).unwrap();
516        assert!(!replica.is_healthy);
517
518        // Update non-existent replica should return false
519        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        // Update lag information
529        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        // Both replicas should be healthy initially
543        assert_eq!(manager.get_healthy_replicas().len(), 2);
544
545        // Run health check with very short threshold - should mark all as unhealthy
546        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        // Add replicas
566        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        // All healthy - should meet requirements
571        assert!(manager.is_replication_healthy());
572
573        // Mark one as unhealthy - should still meet requirements (2 out of 3)
574        manager.update_replica_health(4, false);
575        assert!(manager.is_replication_healthy());
576
577        // Mark another as unhealthy - should not meet requirements (1 out of 3)
578        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        // Synchronous requires all replicas
589        assert_eq!(manager.required_replica_count(), 2);
590
591        // Asynchronous requires no replicas
592        manager.set_strategy(ReplicationStrategy::Asynchronous);
593        assert_eq!(manager.required_replica_count(), 0);
594
595        // Semi-synchronous requires minimum specified
596        manager.set_strategy(ReplicationStrategy::SemiSynchronous { min_replicas: 1 });
597        assert_eq!(manager.required_replica_count(), 1);
598
599        // Raft consensus requires majority
600        manager.set_strategy(ReplicationStrategy::RaftConsensus);
601        // With 2 replicas + 1 local = 3 total, majority is 2, so we need 1 replica (since local is always available)
602        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        // Single node cluster (1 local + 0 replicas = 1 total) - should be healthy
610        assert!(manager.is_replication_healthy());
611
612        // Add replicas to form 3-node cluster
613        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        // All healthy (3 total, majority = 2, local + 2 replicas = 3) - should be healthy
617        assert!(manager.is_replication_healthy());
618
619        // Mark one replica as unhealthy (local + 1 replica = 2, still majority) - should be healthy
620        manager.update_replica_health(3, false);
621        assert!(manager.is_replication_healthy());
622
623        // Mark both replicas as unhealthy (only local = 1, not majority) - should not be healthy
624        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}