ruvector_graph/distributed/
replication.rs1use crate::distributed::shard::{EdgeData, GraphShard, NodeData, NodeId, ShardId};
10use crate::{GraphError, Result};
11use chrono::{DateTime, Utc};
12use dashmap::DashMap;
13use ruvector_replication::{
14 Replica, ReplicaRole, ReplicaSet, ReplicationLog, SyncManager, SyncMode,
15};
16use serde::{Deserialize, Serialize};
17use std::collections::{HashMap, HashSet};
18use std::sync::Arc;
19use tracing::{debug, info, warn};
20use uuid::Uuid;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
24pub enum ReplicationStrategy {
25 FullShard,
27 VertexCut,
29 Subgraph,
31 Hybrid,
33}
34
35#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct GraphReplicationConfig {
38 pub replication_factor: usize,
40 pub strategy: ReplicationStrategy,
42 pub high_degree_threshold: usize,
44 pub sync_mode: SyncMode,
46 pub enable_conflict_resolution: bool,
48 pub timeout_seconds: u64,
50}
51
52impl Default for GraphReplicationConfig {
53 fn default() -> Self {
54 Self {
55 replication_factor: 3,
56 strategy: ReplicationStrategy::FullShard,
57 high_degree_threshold: 100,
58 sync_mode: SyncMode::Async,
59 enable_conflict_resolution: true,
60 timeout_seconds: 30,
61 }
62 }
63}
64
65pub struct GraphReplication {
67 config: GraphReplicationConfig,
69 replica_sets: Arc<DashMap<ShardId, Arc<ReplicaSet>>>,
71 sync_managers: Arc<DashMap<ShardId, Arc<SyncManager>>>,
73 high_degree_nodes: Arc<DashMap<NodeId, usize>>,
75 node_replicas: Arc<DashMap<NodeId, Vec<String>>>,
77}
78
79impl GraphReplication {
80 pub fn new(config: GraphReplicationConfig) -> Self {
82 Self {
83 config,
84 replica_sets: Arc::new(DashMap::new()),
85 sync_managers: Arc::new(DashMap::new()),
86 high_degree_nodes: Arc::new(DashMap::new()),
87 node_replicas: Arc::new(DashMap::new()),
88 }
89 }
90
91 pub fn initialize_shard_replication(
93 &self,
94 shard_id: ShardId,
95 primary_node: String,
96 replica_nodes: Vec<String>,
97 ) -> Result<()> {
98 info!(
99 "Initializing replication for shard {} with {} replicas",
100 shard_id,
101 replica_nodes.len()
102 );
103
104 let mut replica_set = ReplicaSet::new(format!("shard-{}", shard_id));
106
107 replica_set
109 .add_replica(
110 &primary_node,
111 &format!("{}:9001", primary_node),
112 ReplicaRole::Primary,
113 )
114 .map_err(|e| GraphError::ReplicationError(e.to_string()))?;
115
116 for (idx, node) in replica_nodes.iter().enumerate() {
118 replica_set
119 .add_replica(
120 &format!("{}-replica-{}", node, idx),
121 &format!("{}:9001", node),
122 ReplicaRole::Secondary,
123 )
124 .map_err(|e| GraphError::ReplicationError(e.to_string()))?;
125 }
126
127 let replica_set = Arc::new(replica_set);
128
129 let log = Arc::new(ReplicationLog::new(&primary_node));
131
132 let sync_manager = Arc::new(SyncManager::new(Arc::clone(&replica_set), log));
134 sync_manager.set_sync_mode(self.config.sync_mode.clone());
135
136 self.replica_sets.insert(shard_id, replica_set);
137 self.sync_managers.insert(shard_id, sync_manager);
138
139 Ok(())
140 }
141
142 pub async fn replicate_node_add(&self, shard_id: ShardId, node: NodeData) -> Result<()> {
144 debug!(
145 "Replicating node addition: {} to shard {}",
146 node.id, shard_id
147 );
148
149 match self.config.strategy {
151 ReplicationStrategy::FullShard => {
152 self.replicate_to_shard(shard_id, ReplicationOp::AddNode(node))
153 .await
154 }
155 ReplicationStrategy::VertexCut => {
156 let degree = self.get_node_degree(&node.id);
158 if degree >= self.config.high_degree_threshold {
159 self.replicate_high_degree_node(node).await
161 } else {
162 self.replicate_to_shard(shard_id, ReplicationOp::AddNode(node))
163 .await
164 }
165 }
166 ReplicationStrategy::Subgraph | ReplicationStrategy::Hybrid => {
167 self.replicate_to_shard(shard_id, ReplicationOp::AddNode(node))
168 .await
169 }
170 }
171 }
172
173 pub async fn replicate_edge_add(&self, shard_id: ShardId, edge: EdgeData) -> Result<()> {
175 debug!(
176 "Replicating edge addition: {} to shard {}",
177 edge.id, shard_id
178 );
179
180 self.increment_node_degree(&edge.from);
182 self.increment_node_degree(&edge.to);
183
184 self.replicate_to_shard(shard_id, ReplicationOp::AddEdge(edge))
185 .await
186 }
187
188 pub async fn replicate_node_delete(&self, shard_id: ShardId, node_id: NodeId) -> Result<()> {
190 debug!(
191 "Replicating node deletion: {} from shard {}",
192 node_id, shard_id
193 );
194
195 self.replicate_to_shard(shard_id, ReplicationOp::DeleteNode(node_id))
196 .await
197 }
198
199 pub async fn replicate_edge_delete(&self, shard_id: ShardId, edge_id: String) -> Result<()> {
201 debug!(
202 "Replicating edge deletion: {} from shard {}",
203 edge_id, shard_id
204 );
205
206 self.replicate_to_shard(shard_id, ReplicationOp::DeleteEdge(edge_id))
207 .await
208 }
209
210 async fn replicate_to_shard(&self, shard_id: ShardId, op: ReplicationOp) -> Result<()> {
212 let sync_manager = self
213 .sync_managers
214 .get(&shard_id)
215 .ok_or_else(|| GraphError::ShardError(format!("Shard {} not initialized", shard_id)))?;
216
217 let _data =
222 serde_json::to_vec(&op).map_err(|e| GraphError::SerializationError(e.to_string()))?;
223
224 debug!("Replicating operation for shard {}", shard_id);
228
229 Ok(())
230 }
231
232 async fn replicate_high_degree_node(&self, node: NodeData) -> Result<()> {
234 info!(
235 "Replicating high-degree node {} to multiple shards",
236 node.id
237 );
238
239 let degree = self.get_node_degree(&node.id);
241 let replica_count =
242 (degree / self.config.high_degree_threshold).min(self.config.replication_factor);
243
244 let mut replica_shards = Vec::new();
245
246 for shard_id in 0..replica_count {
248 replica_shards.push(shard_id as ShardId);
249 }
250
251 for shard_id in replica_shards.clone() {
253 self.replicate_to_shard(shard_id, ReplicationOp::AddNode(node.clone()))
254 .await?;
255 }
256
257 self.node_replicas.insert(
259 node.id.clone(),
260 replica_shards.iter().map(|s| s.to_string()).collect(),
261 );
262
263 Ok(())
264 }
265
266 fn get_node_degree(&self, node_id: &NodeId) -> usize {
268 self.high_degree_nodes
269 .get(node_id)
270 .map(|d| *d.value())
271 .unwrap_or(0)
272 }
273
274 fn increment_node_degree(&self, node_id: &NodeId) {
276 self.high_degree_nodes
277 .entry(node_id.clone())
278 .and_modify(|d| *d += 1)
279 .or_insert(1);
280 }
281
282 pub fn get_replica_set(&self, shard_id: ShardId) -> Option<Arc<ReplicaSet>> {
284 self.replica_sets
285 .get(&shard_id)
286 .map(|r| Arc::clone(r.value()))
287 }
288
289 pub fn get_sync_manager(&self, shard_id: ShardId) -> Option<Arc<SyncManager>> {
291 self.sync_managers
292 .get(&shard_id)
293 .map(|s| Arc::clone(s.value()))
294 }
295
296 pub fn get_stats(&self) -> ReplicationStats {
298 ReplicationStats {
299 total_shards: self.replica_sets.len(),
300 high_degree_nodes: self.high_degree_nodes.len(),
301 replicated_nodes: self.node_replicas.len(),
302 strategy: self.config.strategy,
303 }
304 }
305
306 pub async fn health_check(&self) -> HashMap<ShardId, ReplicaHealth> {
308 let mut health = HashMap::new();
309
310 for entry in self.replica_sets.iter() {
311 let shard_id = *entry.key();
312 let replica_set = entry.value();
313
314 let healthy_count = self.config.replication_factor;
316
317 health.insert(
318 shard_id,
319 ReplicaHealth {
320 total_replicas: self.config.replication_factor,
321 healthy_replicas: healthy_count,
322 is_healthy: healthy_count >= (self.config.replication_factor / 2 + 1),
323 },
324 );
325 }
326
327 health
328 }
329
330 pub fn config(&self) -> &GraphReplicationConfig {
332 &self.config
333 }
334}
335
336#[derive(Debug, Clone, Serialize, Deserialize)]
338enum ReplicationOp {
339 AddNode(NodeData),
340 AddEdge(EdgeData),
341 DeleteNode(NodeId),
342 DeleteEdge(String),
343 UpdateNode(NodeData),
344 UpdateEdge(EdgeData),
345}
346
347#[derive(Debug, Clone, Serialize, Deserialize)]
349pub struct ReplicationStats {
350 pub total_shards: usize,
351 pub high_degree_nodes: usize,
352 pub replicated_nodes: usize,
353 pub strategy: ReplicationStrategy,
354}
355
356#[derive(Debug, Clone, Serialize, Deserialize)]
358pub struct ReplicaHealth {
359 pub total_replicas: usize,
360 pub healthy_replicas: usize,
361 pub is_healthy: bool,
362}
363
364#[cfg(test)]
365mod tests {
366 use super::*;
367 use std::collections::HashMap;
368
369 #[tokio::test]
370 async fn test_graph_replication() {
371 let config = GraphReplicationConfig::default();
372 let replication = GraphReplication::new(config);
373
374 replication
375 .initialize_shard_replication(0, "node-1".to_string(), vec!["node-2".to_string()])
376 .unwrap();
377
378 assert!(replication.get_replica_set(0).is_some());
379 assert!(replication.get_sync_manager(0).is_some());
380 }
381
382 #[tokio::test]
383 async fn test_node_replication() {
384 let config = GraphReplicationConfig::default();
385 let replication = GraphReplication::new(config);
386
387 replication
388 .initialize_shard_replication(0, "node-1".to_string(), vec!["node-2".to_string()])
389 .unwrap();
390
391 let node = NodeData {
392 id: "test-node".to_string(),
393 properties: HashMap::new(),
394 labels: vec!["Test".to_string()],
395 };
396
397 let result = replication.replicate_node_add(0, node).await;
398 assert!(result.is_ok());
399 }
400
401 #[test]
402 fn test_replication_stats() {
403 let config = GraphReplicationConfig::default();
404 let replication = GraphReplication::new(config);
405
406 let stats = replication.get_stats();
407 assert_eq!(stats.total_shards, 0);
408 assert_eq!(stats.strategy, ReplicationStrategy::FullShard);
409 }
410
411 #[test]
412 fn replication_operations_preserve_dynamic_json_properties() {
413 let op = ReplicationOp::AddNode(NodeData {
414 id: "node-1".to_string(),
415 properties: HashMap::from([(
416 "kind".to_string(),
417 serde_json::Value::String("test".to_string()),
418 )]),
419 labels: vec!["Example".to_string()],
420 });
421
422 let encoded = serde_json::to_vec(&op).unwrap();
423 let decoded: ReplicationOp = serde_json::from_slice(&encoded).unwrap();
424 match decoded {
425 ReplicationOp::AddNode(node) => {
426 assert_eq!(node.id, "node-1");
427 assert_eq!(node.labels, vec!["Example"]);
428 }
429 _ => panic!("unexpected replication operation"),
430 }
431 }
432}