shadow_load_testing/
swarm.rs

1//! Virtual peer swarm for load testing
2
3use shadow_core::{PeerId, PeerInfo};
4use dht::RoutingTable;
5
6/// A virtual swarm of simulated peers
7pub struct VirtualSwarm {
8    peers: Vec<VirtualPeer>,
9}
10
11/// A single virtual peer with its own ID and routing table
12pub struct VirtualPeer {
13    id: PeerId,
14    routing_table: RoutingTable,
15}
16
17impl VirtualSwarm {
18    /// Create a swarm of N virtual peers
19    pub fn new(count: usize) -> Self {
20        let peers: Vec<VirtualPeer> = (0..count)
21            .map(|i| {
22                let id = PeerId::random();
23                let routing_table = RoutingTable::new(id, 20);
24                VirtualPeer { id, routing_table }
25            })
26            .collect();
27
28        VirtualSwarm { peers }
29    }
30
31    /// Get the number of peers
32    pub fn peer_count(&self) -> usize {
33        self.peers.len()
34    }
35
36    /// Get all peer IDs
37    pub fn peer_ids(&self) -> Vec<PeerId> {
38        self.peers.iter().map(|p| p.id).collect()
39    }
40
41    /// Get a reference to a specific peer
42    pub fn peer(&self, index: usize) -> Option<&VirtualPeer> {
43        self.peers.get(index)
44    }
45
46    /// Get a mutable reference to a specific peer
47    pub fn peer_mut(&mut self, index: usize) -> Option<&mut VirtualPeer> {
48        self.peers.get_mut(index)
49    }
50
51    /// Bootstrap the swarm: each peer learns about some neighbors
52    pub fn bootstrap(&mut self, connections_per_peer: usize) {
53        let ids: Vec<PeerId> = self.peer_ids();
54        let count = ids.len();
55
56        for i in 0..count {
57            let step = std::cmp::max(1, count / connections_per_peer);
58            for k in 1..=connections_per_peer {
59                let j = (i + k * step) % count;
60                if j != i {
61                    let peer_info = PeerInfo::new(
62                        ids[j],
63                        vec![format!("127.0.0.1:{}", 10000 + j)],
64                        [0u8; 32],
65                        [0u8; 32],
66                    );
67                    let _ = self.peers[i].routing_table.add_peer(peer_info);
68                }
69            }
70        }
71    }
72
73    /// Get average routing table size across all peers
74    pub fn avg_routing_table_size(&self) -> f64 {
75        if self.peers.is_empty() {
76            return 0.0;
77        }
78        let total: usize = self.peers.iter()
79            .map(|p| p.routing_table.peer_count())
80            .sum();
81        total as f64 / self.peers.len() as f64
82    }
83
84    /// Find peers closest to a target across the whole swarm
85    pub fn find_closest_global(&self, target: &PeerId, k: usize) -> Vec<PeerId> {
86        let mut all_ids = self.peer_ids();
87        all_ids.sort_by(|a, b| {
88            let dist_a = a.xor_distance(target);
89            let dist_b = b.xor_distance(target);
90            dist_a.cmp(&dist_b)
91        });
92        all_ids.truncate(k);
93        all_ids
94    }
95
96    /// Simulate iterative lookup from a starting peer
97    pub fn iterative_lookup(&self, start: usize, target: &PeerId, k: usize) -> Vec<PeerInfo> {
98        if let Some(peer) = self.peers.get(start) {
99            peer.routing_table.find_closest(target, k)
100        } else {
101            vec![]
102        }
103    }
104}
105
106impl VirtualPeer {
107    pub fn id(&self) -> PeerId {
108        self.id
109    }
110
111    pub fn routing_table(&self) -> &RoutingTable {
112        &self.routing_table
113    }
114
115    pub fn routing_table_mut(&mut self) -> &mut RoutingTable {
116        &mut self.routing_table
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn test_swarm_creation() {
126        let swarm = VirtualSwarm::new(100);
127        assert_eq!(swarm.peer_count(), 100);
128        assert_eq!(swarm.peer_ids().len(), 100);
129    }
130
131    #[test]
132    fn test_swarm_bootstrap() {
133        let mut swarm = VirtualSwarm::new(50);
134        swarm.bootstrap(5);
135        let avg = swarm.avg_routing_table_size();
136        assert!(avg > 0.0, "Should have peers in routing tables after bootstrap");
137    }
138
139    #[test]
140    fn test_swarm_lookup() {
141        let mut swarm = VirtualSwarm::new(100);
142        swarm.bootstrap(10);
143        let target = PeerId::random();
144        let closest = swarm.find_closest_global(&target, 20);
145        assert_eq!(closest.len(), 20);
146    }
147
148    #[test]
149    fn test_unique_peer_ids() {
150        let swarm = VirtualSwarm::new(1000);
151        let ids = swarm.peer_ids();
152        let unique: std::collections::HashSet<_> = ids.iter().collect();
153        assert_eq!(unique.len(), 1000, "All peer IDs should be unique");
154    }
155
156    #[test]
157    fn test_large_swarm() {
158        let swarm = VirtualSwarm::new(10_000);
159        assert_eq!(swarm.peer_count(), 10_000);
160    }
161
162    #[test]
163    fn test_iterative_lookup() {
164        let mut swarm = VirtualSwarm::new(50);
165        swarm.bootstrap(5);
166        let target = PeerId::random();
167        let results = swarm.iterative_lookup(0, &target, 5);
168        // Should find some peers (depends on bootstrap)
169        assert!(!results.is_empty() || swarm.peer(0).unwrap().routing_table().peer_count() == 0);
170    }
171}
shadow_load_testing/swarm.rs

shadow_load_testing/
swarm.rs
