File size: 5,074 Bytes
d8435ba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
use api::grpc::qdrant::raft_server::Raft;
use api::grpc::qdrant::{
    AddPeerToKnownMessage, AllPeers, Peer, PeerId, RaftMessage as RaftMessageBytes, Uri as UriStr,
};
use itertools::Itertools;
use raft::eraftpb::Message as RaftMessage;
use storage::content_manager::consensus_manager::ConsensusStateRef;
use storage::content_manager::consensus_ops::ConsensusOperations;
use tokio::sync::mpsc::Sender;
use tonic::transport::Uri;
use tonic::{async_trait, Request, Response, Status};

use super::validate;
use crate::consensus;

pub struct RaftService {
    message_sender: Sender<consensus::Message>,
    consensus_state: ConsensusStateRef,
}

impl RaftService {
    pub fn new(sender: Sender<consensus::Message>, consensus_state: ConsensusStateRef) -> Self {
        Self {
            message_sender: sender,
            consensus_state,
        }
    }
}

#[async_trait]
impl Raft for RaftService {
    async fn send(&self, mut request: Request<RaftMessageBytes>) -> Result<Response<()>, Status> {
        let message =
            <RaftMessage as prost_for_raft::Message>::decode(&request.get_mut().message[..])
                .map_err(|err| {
                    Status::invalid_argument(format!("Failed to parse raft message: {err}"))
                })?;
        self.message_sender
            .send(consensus::Message::FromPeer(Box::new(message)))
            .await
            .map_err(|_| Status::internal("Can't send Raft message over channel"))?;
        Ok(Response::new(()))
    }

    async fn who_is(
        &self,
        request: tonic::Request<PeerId>,
    ) -> Result<tonic::Response<UriStr>, tonic::Status> {
        let addresses = self.consensus_state.peer_address_by_id();
        let uri = addresses
            .get(&request.get_ref().id)
            .ok_or_else(|| Status::internal("Peer not found"))?;
        Ok(Response::new(UriStr {
            uri: uri.to_string(),
        }))
    }

    async fn add_peer_to_known(
        &self,
        request: tonic::Request<AddPeerToKnownMessage>,
    ) -> Result<tonic::Response<AllPeers>, tonic::Status> {
        validate(request.get_ref())?;
        let peer = request.get_ref();
        let uri_string = if let Some(uri) = &peer.uri {
            uri.clone()
        } else {
            let ip = request
                .remote_addr()
                .ok_or_else(|| {
                    Status::failed_precondition("Remote address unavailable due to the used IO")
                })?
                .ip();
            let port = peer
                .port
                .ok_or_else(|| Status::invalid_argument("URI or port should be supplied"))?;
            format!("http://{ip}:{port}")
        };
        let uri: Uri = uri_string
            .parse()
            .map_err(|err| Status::internal(format!("Failed to parse uri: {err}")))?;
        let peer = request.into_inner();

        // the consensus operation can take up to DEFAULT_META_OP_WAIT
        self.consensus_state
            .propose_consensus_op_with_await(
                ConsensusOperations::AddPeer {
                    peer_id: peer.id,
                    uri: uri.to_string(),
                },
                None,
            )
            .await
            .map_err(|err| Status::internal(format!("Failed to add peer: {err}")))?;

        let mut addresses = self.consensus_state.peer_address_by_id();

        // Make sure that the new peer is now present in the known addresses
        if !addresses.values().contains(&uri) {
            return Err(Status::internal(format!(
                "Failed to add peer after consensus: {uri}"
            )));
        }

        let first_peer_id = self.consensus_state.first_voter();

        // If `first_peer_id` is not present in the list of peers, it means it was removed from
        // cluster at some point.
        //
        // Before Qdrant version 1.11.6 origin peer was not committed to consensus, so if it was
        // removed from cluster, any node added to the cluster after this would not recognize it as
        // being part of the cluster in the past and will end up with a broken consensus state.
        //
        // To prevent this, we add `first_peer_id` (with a fake URI) to the list of peers.
        //
        // `add_peer_to_known` is used to add new peers to the cluster, and so `first_peer_id` (and
        // its fake URI) would be removed from new peer's state shortly, while it will be synchronizing
        // and applying past Raft log.
        addresses.entry(first_peer_id).or_default();

        Ok(Response::new(AllPeers {
            all_peers: addresses
                .into_iter()
                .map(|(id, uri)| Peer {
                    id,
                    uri: uri.to_string(),
                })
                .collect(),
            first_peer_id,
        }))
    }

    // Left for compatibility - does nothing
    async fn add_peer_as_participant(
        &self,
        _request: tonic::Request<PeerId>,
    ) -> Result<tonic::Response<()>, tonic::Status> {
        Ok(Response::new(()))
    }
}