ruCCL 0.21.1

Ruda collective communication algorithms and orchestration.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use std::{
    marker::PhantomData,
    sync::{Arc, Mutex},
    vec,
};

use ruda_communication::{CommunicationChannel, Message, Protocol, ProtocolClient};
use serde::{Deserialize, Serialize};
use tokio::sync::{Notify, RwLock};

use crate::{NodeId, node::base::NodeState};

/// Handles the status of sync requests from other nodes
pub(crate) struct SyncService<P: Protocol> {
    /// Current node's state, shared with the thread that does aggregations
    node_state: Arc<RwLock<Option<NodeState>>>,
    /// The number of peers that have requested to sync with us since the last successful sync.
    syncing_peers: Mutex<Vec<NodeId>>,
    /// Notification on each incoming sync request
    sync_notif: Notify,

    _p: PhantomData<P>,
}

#[derive(Debug, Serialize, Deserialize)]
struct SyncRequest(NodeId);

impl<P: Protocol> SyncService<P> {
    pub fn new(node_state: Arc<RwLock<Option<NodeState>>>) -> Self {
        Self {
            node_state,
            syncing_peers: Mutex::new(vec![]),
            sync_notif: Notify::new(),
            _p: PhantomData,
        }
    }

    fn add_syncing_peer(&self, peer: NodeId) {
        let mut syncing_peers = self.syncing_peers.lock().unwrap();
        syncing_peers.push(peer);
    }

    /// Sync with all peers.
    pub async fn sync(&self) {
        // we can't sync while we register
        let node_state = self.node_state.read().await;
        let node_state = node_state
            .as_ref()
            .expect("Trying to sync a node before having registered to the orchestrator");

        // this peer is syncing
        self.add_syncing_peer(node_state.node_id);
        for (id, addr) in &node_state.nodes {
            if *id == node_state.node_id {
                continue;
            }

            let mut connection = P::Client::connect(addr.clone(), "sync")
                .await
                .expect("Couldn't connect to peer for sync");
            let msg = SyncRequest(node_state.node_id);
            let sync_bytes = rmp_serde::to_vec(&msg).unwrap();
            connection
                .send(Message::new(sync_bytes.into()))
                .await
                .expect("Peer closed connection unexpectedly");
        }
        loop {
            {
                // compare currently synced peers with list of all nodes
                let mut syncing_peers = self.syncing_peers.lock().unwrap().to_vec();
                syncing_peers.sort();

                let mut all_node_ids = node_state.nodes.keys().cloned().collect::<Vec<_>>();
                all_node_ids.sort();

                if syncing_peers == all_node_ids {
                    // all nodes have synced
                    syncing_peers.clear();
                    return;
                }
            }
            // Wait for the next sync to come in
            self.sync_notif.notified().await
        }
    }

    pub async fn handle_sync_connection<C: CommunicationChannel>(&self, mut channel: C) {
        let msg = channel.recv().await.unwrap();
        let Some(msg) = msg else {
            return;
        };

        let msg = rmp_serde::from_slice::<SyncRequest>(&msg.data).unwrap();

        self.add_syncing_peer(msg.0);

        self.sync_notif.notify_waiters();
    }
}