sui_network/discovery/
server.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use super::{Discovery, MAX_PEERS_TO_SEND, SignedNodeInfo, State};
5use anemo::{Request, Response};
6use rand::seq::IteratorRandom;
7use serde::{Deserialize, Serialize};
8use std::{
9    collections::HashMap,
10    sync::{Arc, RwLock},
11};
12
13#[derive(Clone, Debug, Serialize, Deserialize)]
14pub struct GetKnownPeersResponseV2 {
15    pub own_info: SignedNodeInfo,
16    pub known_peers: Vec<SignedNodeInfo>,
17}
18
19pub(super) struct Server {
20    pub(super) state: Arc<RwLock<State>>,
21}
22
23#[anemo::async_trait]
24impl Discovery for Server {
25    async fn get_known_peers_v2(
26        &self,
27        _request: Request<()>,
28    ) -> Result<Response<GetKnownPeersResponseV2>, anemo::rpc::Status> {
29        let state = self.state.read().unwrap();
30        let own_info = state
31            .our_info
32            .clone()
33            .ok_or_else(|| anemo::rpc::Status::internal("own_info has not been initialized yet"))?;
34
35        let known_peers = if state.known_peers.len() < MAX_PEERS_TO_SEND {
36            state
37                .known_peers
38                .values()
39                .map(|e| e.inner())
40                .cloned()
41                .collect()
42        } else {
43            let mut rng = rand::thread_rng();
44            // prefer returning peers that we are connected to as they are known-good
45            let mut known_peers = state
46                .connected_peers
47                .keys()
48                .filter_map(|peer_id| state.known_peers.get(peer_id))
49                .map(|info| (info.peer_id, info))
50                .choose_multiple(&mut rng, MAX_PEERS_TO_SEND)
51                .into_iter()
52                .collect::<HashMap<_, _>>();
53
54            if known_peers.len() <= MAX_PEERS_TO_SEND {
55                // Fill the remaining space with other peers, randomly sampling at most MAX_PEERS_TO_SEND
56                for info in state
57                    .known_peers
58                    .values()
59                    // This randomly samples the iterator stream but the order of elements after
60                    // sampling may not be random, this is ok though since we're just trying to do
61                    // best-effort on sharing info of peers we haven't connected with ourselves.
62                    .choose_multiple(&mut rng, MAX_PEERS_TO_SEND)
63                {
64                    if known_peers.len() >= MAX_PEERS_TO_SEND {
65                        break;
66                    }
67
68                    known_peers.insert(info.peer_id, info);
69                }
70            }
71
72            known_peers
73                .into_values()
74                .map(|e| e.inner())
75                .cloned()
76                .collect()
77        };
78
79        Ok(Response::new(GetKnownPeersResponseV2 {
80            own_info,
81            known_peers,
82        }))
83    }
84}
sui_network/discovery/server.rs

sui_network/discovery/
server.rs
