1use async_trait::async_trait;
2
3use crate::dht::entry::PlacedEntry;
4use crate::dht::Chord;
5use crate::dht::ChordStorageSync;
6use crate::dht::Did;
7use crate::dht::PeerRingAction;
8use crate::dht::PeerRingRemoteAction;
9use crate::error::Error;
10use crate::error::Result;
11use crate::message::effects::ConnectionFunctor;
12use crate::message::effects::MessageSendFunctor;
13use crate::message::effects::PayloadRelayFunctor;
14use crate::message::types::Message;
15use crate::message::types::NotifyPredecessorReport;
16use crate::message::types::NotifyPredecessorSend;
17use crate::message::types::SyncEntriesWithSuccessor;
18use crate::message::HandleMsg;
19use crate::message::MessageHandler;
20use crate::message::MessagePayload;
21
22fn collect_sync_entries_actions(
23 act: PeerRingAction,
24 out: &mut Vec<(Did, Vec<PlacedEntry>)>,
25) -> Result<()> {
26 match act {
27 PeerRingAction::None => Ok(()),
28 PeerRingAction::RemoteAction(
29 next,
30 PeerRingRemoteAction::SyncEntriesWithSuccessor(data),
31 ) => {
32 out.push((next, data));
33 Ok(())
34 }
35 PeerRingAction::MultiActions(actions) => {
36 for action in actions {
37 collect_sync_entries_actions(action, out)?;
38 }
39 Ok(())
40 }
41 action => Err(Error::unexpected_peer_ring_action(action)),
42 }
43}
44
45#[cfg_attr(feature = "wasm", async_trait(?Send))]
46#[cfg_attr(not(feature = "wasm"), async_trait)]
47impl HandleMsg<NotifyPredecessorSend> for MessageHandler {
48 async fn handle(&self, ctx: &MessagePayload, msg: &NotifyPredecessorSend) -> Result<()> {
49 let predecessor = self.dht.notify(msg.did)?;
50
51 if predecessor != ctx.relay.try_origin_sender()? {
52 return self
53 .run_effects([PayloadRelayFunctor::send_report_message(
54 ctx,
55 Message::NotifyPredecessorReport(NotifyPredecessorReport { did: predecessor }),
56 )
57 .into()])
58 .await;
59 }
60
61 Ok(())
62 }
63}
64
65#[cfg_attr(feature = "wasm", async_trait(?Send))]
66#[cfg_attr(not(feature = "wasm"), async_trait)]
67impl HandleMsg<NotifyPredecessorReport> for MessageHandler {
68 async fn handle(&self, _ctx: &MessagePayload, msg: &NotifyPredecessorReport) -> Result<()> {
69 self.run_effects([ConnectionFunctor::connect_dht_peer(msg.did).into()])
70 .await?;
71
72 let mut sync_actions = Vec::new();
73 collect_sync_entries_actions(
74 self.dht.sync_entries_with_successor(msg.did).await?,
75 &mut sync_actions,
76 )?;
77 let effects = sync_actions
78 .into_iter()
79 .map(|(next, data)| {
80 MessageSendFunctor::send_message(
81 Message::SyncEntriesWithSuccessor(SyncEntriesWithSuccessor { data }),
82 next,
83 )
84 .into()
85 })
86 .collect::<Vec<_>>();
87 self.run_effects(effects).await?;
88
89 Ok(())
90 }
91}
92
93#[cfg(not(feature = "wasm"))]
94#[cfg(test)]
95mod test {
96 use std::sync::Arc;
97
98 use tokio::time::timeout;
99 use tokio::time::Duration;
100
101 use super::*;
102 use crate::dht::entry::Entry;
103 use crate::dht::entry::EntryKind;
104 use crate::dht::entry::PlacedEntry;
105 use crate::dht::entry::SyncedEntryAck;
106 use crate::dht::successor::SuccessorReader;
107 use crate::ecc::tests::gen_ordered_keys;
108 use crate::ecc::SecretKey;
109 use crate::error::Error;
110 use crate::message::Encoder;
111 use crate::message::SyncEntriesWithSuccessorReport;
112 use crate::session::SessionSk;
113 use crate::swarm::callback::SwarmCallback;
114 use crate::swarm::Swarm;
115 use crate::tests::default::assert_no_more_msg;
116 use crate::tests::default::prepare_node;
117 use crate::tests::default::wait_for_msgs;
118 use crate::tests::manually_establish_connection;
119
120 struct NoopCallback;
121
122 impl SwarmCallback for NoopCallback {}
123
124 fn next_generated_key(keys: &mut impl Iterator<Item = SecretKey>) -> Result<SecretKey> {
125 keys.next()
126 .ok_or_else(|| Error::InvalidMessage("expected generated key".to_string()))
127 }
128
129 #[tokio::test]
130 async fn test_triple_nodes_stabilization_1_2_3() -> Result<()> {
131 let keys = gen_ordered_keys(3);
132 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
133 test_triple_ordered_nodes_stabilization(key1, key2, key3).await
134 }
135
136 #[tokio::test]
137 async fn test_triple_nodes_stabilization_2_3_1() -> Result<()> {
138 let keys = gen_ordered_keys(3);
139 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
140
141 test_triple_ordered_nodes_stabilization(key2, key3, key1).await
142 }
143
144 #[tokio::test]
145 async fn test_triple_nodes_stabilization_3_1_2() -> Result<()> {
146 let keys = gen_ordered_keys(3);
147 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
148 test_triple_ordered_nodes_stabilization(key3, key1, key2).await
149 }
150
151 #[tokio::test]
152 async fn test_triple_nodes_stabilization_3_2_1() -> Result<()> {
153 let keys = gen_ordered_keys(3);
154 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
155 test_triple_desc_ordered_nodes_stabilization(key3, key2, key1).await
156 }
157
158 #[tokio::test]
159 async fn test_triple_nodes_stabilization_2_1_3() -> Result<()> {
160 let keys = gen_ordered_keys(3);
161 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
162 test_triple_desc_ordered_nodes_stabilization(key2, key1, key3).await
163 }
164
165 #[tokio::test]
166 async fn test_triple_nodes_stabilization_1_3_2() -> Result<()> {
167 let keys = gen_ordered_keys(3);
168 let (key1, key2, key3) = (keys[0], keys[1], keys[2]);
169 test_triple_desc_ordered_nodes_stabilization(key1, key3, key2).await
170 }
171
172 #[tokio::test]
173 async fn notify_predecessor_report_syncs_entries_when_predecessor_already_connected(
174 ) -> Result<()> {
175 let mut keys = gen_ordered_keys(3).into_iter();
176 let key1 = next_generated_key(&mut keys)?;
177 let key2 = next_generated_key(&mut keys)?;
178 let key3 = next_generated_key(&mut keys)?;
179
180 let node1 = prepare_node(key1).await;
181 let node2 = prepare_node(key2).await;
182 manually_establish_connection(&node1.swarm, &node2.swarm).await;
183 wait_for_msgs([&node1, &node2]).await;
184 assert_no_more_msg([&node1, &node2]).await;
185
186 let entry = Entry::new(
187 key3.address().into(),
188 vec![String::from("sync me").encode()?],
189 EntryKind::Data,
190 );
191 let stored_entry = entry.clone().try_into_storage_entry()?;
192 node1
193 .dht()
194 .storage
195 .put(&entry.did.to_string(), &entry)
196 .await?;
197
198 let context_key = SecretKey::random();
199 let context_session = SessionSk::new_with_seckey(&context_key)?;
200 let context = MessagePayload::new_send(
201 Message::custom(b"notify report context")?,
202 &context_session,
203 node1.did(),
204 node1.did(),
205 )?;
206
207 let handler = MessageHandler::new(node1.swarm.transport.clone(), Arc::new(NoopCallback));
208 handler
209 .handle(&context, &NotifyPredecessorReport { did: node2.did() })
210 .await?;
211
212 let payload = match timeout(Duration::from_secs(1), node2.listen_once()).await {
213 Ok(Some(payload)) => payload,
214 Ok(None) => {
215 return Err(Error::InvalidMessage(
216 "node2 message stream closed before entry sync".to_string(),
217 ))
218 }
219 Err(_) => {
220 return Err(Error::InvalidMessage(
221 "timed out waiting for entry sync".to_string(),
222 ))
223 }
224 };
225
226 match payload.transaction.data::<Message>()? {
227 Message::SyncEntriesWithSuccessor(SyncEntriesWithSuccessor { data }) => {
228 assert_eq!(data, vec![PlacedEntry::new(entry.did, entry.clone())]);
229 }
230 message => {
231 return Err(Error::InvalidMessage(format!(
232 "expected SyncEntriesWithSuccessor, got {message:?}"
233 )))
234 }
235 }
236 let payload = match timeout(Duration::from_secs(1), node1.listen_once()).await {
237 Ok(Some(payload)) => payload,
238 Ok(None) => {
239 return Err(Error::InvalidMessage(
240 "node1 message stream closed before entry sync ack".to_string(),
241 ))
242 }
243 Err(_) => {
244 return Err(Error::InvalidMessage(
245 "timed out waiting for entry sync ack".to_string(),
246 ))
247 }
248 };
249
250 match payload.transaction.data::<Message>()? {
251 Message::SyncEntriesWithSuccessorReport(SyncEntriesWithSuccessorReport { acks }) => {
252 assert_eq!(acks, vec![SyncedEntryAck::new(
253 entry.did,
254 stored_entry.clone()
255 )]);
256 }
257 message => {
258 return Err(Error::InvalidMessage(format!(
259 "expected SyncEntriesWithSuccessorReport, got {message:?}"
260 )))
261 }
262 }
263 assert_eq!(node1.dht().storage.get(&entry.did.to_string()).await?, None);
264 assert_eq!(
265 node2.dht().storage.get(&entry.did.to_string()).await?,
266 Some(stored_entry)
267 );
268
269 Ok(())
270 }
271
272 async fn test_triple_ordered_nodes_stabilization(
273 key1: SecretKey,
274 key2: SecretKey,
275 key3: SecretKey,
276 ) -> Result<()> {
277 let node1 = prepare_node(key1).await;
278 let node2 = prepare_node(key2).await;
279 let node3 = prepare_node(key3).await;
280
281 println!("========================================");
282 println!("|| now we connect node1 and node2 ||");
283 println!("========================================");
284
285 manually_establish_connection(&node1.swarm, &node2.swarm).await;
286 wait_for_msgs([&node1, &node2, &node3]).await;
287 assert_no_more_msg([&node1, &node2, &node3]).await;
288
289 println!("========================================");
290 println!("|| now we start join node3 to node2 ||");
291 println!("========================================");
292
293 manually_establish_connection(&node3.swarm, &node2.swarm).await;
294 wait_for_msgs([&node1, &node2, &node3]).await;
295 assert_no_more_msg([&node1, &node2, &node3]).await;
296
297 println!("=== Check state before stabilization ===");
298 assert_eq!(node1.dht().successors().list()?, vec![node2.did()]);
299 assert_eq!(node2.dht().successors().list()?, vec![
300 node3.did(),
301 node1.did()
302 ]);
303 assert_eq!(node3.dht().successors().list()?, vec![node2.did()]);
304 assert!(node1.dht().lock_predecessor()?.is_none());
305 assert!(node2.dht().lock_predecessor()?.is_none());
306 assert!(node3.dht().lock_predecessor()?.is_none());
307
308 println!("========================================");
309 println!("|| now we start first stabilization ||");
310 println!("========================================");
311
312 run_stabilization_once(node1.swarm.clone()).await?;
313 run_stabilization_once(node2.swarm.clone()).await?;
314 run_stabilization_once(node3.swarm.clone()).await?;
315
316 wait_for_msgs([&node1, &node2, &node3]).await;
317 assert_no_more_msg([&node1, &node2, &node3]).await;
318
319 println!("=== Check state after first stabilization ===");
320 assert!(node1.dht().successors().list()?.contains(&node2.did()));
321 assert_eq!(node2.dht().successors().list()?, vec![
322 node3.did(),
323 node1.did()
324 ]);
325 assert!(node3.dht().successors().list()?.contains(&node2.did()));
326
327 println!("==========================================");
328 println!("|| now we start 5 times stabilization ||");
329 println!("==========================================");
330
331 for _ in 0..5 {
332 run_stabilization_once(node1.swarm.clone()).await?;
333 run_stabilization_once(node2.swarm.clone()).await?;
334 run_stabilization_once(node3.swarm.clone()).await?;
335
336 wait_for_msgs([&node1, &node2, &node3]).await;
337 assert_no_more_msg([&node1, &node2, &node3]).await;
338
339 println!("=== Check state after stabilization ===");
340 assert_eq!(node1.dht().successors().list()?, vec![
341 node2.did(),
342 node3.did()
343 ]);
344 assert_eq!(node2.dht().successors().list()?, vec![
345 node3.did(),
346 node1.did()
347 ]);
348 assert_eq!(node3.dht().successors().list()?, vec![
349 node1.did(),
350 node2.did()
351 ]);
352 }
353
354 println!("=== Check predecessor after all stabilization ===");
355 assert_eq!(*node1.dht().lock_predecessor()?, Some(node3.did()));
356 assert_eq!(*node2.dht().lock_predecessor()?, Some(node1.did()));
357 assert_eq!(*node3.dht().lock_predecessor()?, Some(node2.did()));
358 Ok(())
359 }
360
361 async fn test_triple_desc_ordered_nodes_stabilization(
362 key1: SecretKey,
363 key2: SecretKey,
364 key3: SecretKey,
365 ) -> Result<()> {
366 let node1 = prepare_node(key1).await;
367 let node2 = prepare_node(key2).await;
368 let node3 = prepare_node(key3).await;
369
370 println!("========================================");
371 println!("|| now we connect node1 and node2 ||");
372 println!("========================================");
373
374 manually_establish_connection(&node1.swarm, &node2.swarm).await;
375 wait_for_msgs([&node1, &node2, &node3]).await;
376 assert_no_more_msg([&node1, &node2, &node3]).await;
377
378 println!("========================================");
379 println!("|| now we start join node3 to node2 ||");
380 println!("========================================");
381
382 manually_establish_connection(&node3.swarm, &node2.swarm).await;
383 wait_for_msgs([&node1, &node2, &node3]).await;
384 assert_no_more_msg([&node1, &node2, &node3]).await;
385
386 println!("=== Check state before stabilization ===");
387 assert_eq!(node1.dht().successors().list()?, vec![node2.did()]);
388 assert_eq!(node2.dht().successors().list()?, vec![
389 node1.did(),
390 node3.did()
391 ]);
392 assert_eq!(node3.dht().successors().list()?, vec![node2.did()]);
393 assert!(node1.dht().lock_predecessor()?.is_none());
394 assert!(node2.dht().lock_predecessor()?.is_none());
395 assert!(node3.dht().lock_predecessor()?.is_none());
396
397 println!("========================================");
398 println!("|| now we start first stabilization ||");
399 println!("========================================");
400
401 run_stabilization_once(node1.swarm.clone()).await?;
402 run_stabilization_once(node2.swarm.clone()).await?;
403 run_stabilization_once(node3.swarm.clone()).await?;
404
405 wait_for_msgs([&node1, &node2, &node3]).await;
406 assert_no_more_msg([&node1, &node2, &node3]).await;
407
408 println!("=== Check state after first stabilization ===");
409 assert!(node1.dht().successors().list()?.contains(&node2.did()));
410 assert_eq!(node2.dht().successors().list()?, vec![
411 node1.did(),
412 node3.did()
413 ]);
414 assert!(node3.dht().successors().list()?.contains(&node2.did()));
415
416 println!("==========================================");
417 println!("|| now we start 5 times stabilization ||");
418 println!("==========================================");
419
420 for _ in 0..5 {
421 run_stabilization_once(node1.swarm.clone()).await?;
422 run_stabilization_once(node2.swarm.clone()).await?;
423 run_stabilization_once(node3.swarm.clone()).await?;
424
425 wait_for_msgs([&node1, &node2, &node3]).await;
426 assert_no_more_msg([&node1, &node2, &node3]).await;
427
428 println!("=== Check state after stabilization ===");
429 assert_eq!(node1.dht().successors().list()?, vec![
430 node3.did(),
431 node2.did()
432 ]);
433 assert_eq!(node2.dht().successors().list()?, vec![
434 node1.did(),
435 node3.did()
436 ]);
437 assert_eq!(node3.dht().successors().list()?, vec![
438 node2.did(),
439 node1.did()
440 ]);
441 }
442
443 println!("=== Check predecessor after all stabilization ===");
444 assert_eq!(*node1.dht().lock_predecessor()?, Some(node2.did()));
445 assert_eq!(*node2.dht().lock_predecessor()?, Some(node3.did()));
446 assert_eq!(*node3.dht().lock_predecessor()?, Some(node1.did()));
447
448 Ok(())
449 }
450
451 async fn run_stabilization_once(swarm: Arc<Swarm>) -> Result<()> {
452 let stab = swarm.stabilizer();
453 stab.notify_predecessor().await
454 }
455}