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

sui_network/discovery/
server.rs
