hyperswarm/
peer_discovery.rs

1use std::{
2    net::SocketAddr,
3    pin::Pin,
4    task::{Context, Poll},
5    time::Duration,
6};
7
8use dht_rpc::{Commit, IdBytes};
9use futures::Stream;
10use hyperdht::{
11    Keypair,
12    adht::{Announce, Dht, Lookup, LookupResponse},
13};
14use tokio::time::Sleep;
15
16use crate::JoinOpts;
17
18const REFRESH_BASE_SECS: u64 = 600; // ~10 min
19const REFRESH_JITTER_SECS: u64 = 120; // up to ~2 min
20
21enum ActiveQuery {
22    Lookup(Lookup),
23    Announce(Announce),
24}
25
26impl Stream for ActiveQuery {
27    type Item = hyperdht::Result<Option<LookupResponse>>;
28
29    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
30        match self.get_mut() {
31            ActiveQuery::Lookup(lookup) => Pin::new(lookup).poll_next(cx),
32            ActiveQuery::Announce(announce) => Pin::new(announce).poll_next(cx),
33        }
34    }
35}
36
37enum PeerDiscoveryState {
38    Querying(Box<ActiveQuery>),
39    Sleeping(Pin<Box<Sleep>>),
40}
41
42pub struct PeerDiscovery {
43    topic: IdBytes,
44    opts: JoinOpts,
45    dht: Dht,
46    keypair: Keypair,
47    state: PeerDiscoveryState,
48    first_round_complete: bool,
49    relay_addresses: Vec<SocketAddr>,
50}
51
52impl PeerDiscovery {
53    pub fn new(
54        topic: IdBytes,
55        opts: JoinOpts,
56        dht: Dht,
57        keypair: Keypair,
58        relay_addresses: Vec<SocketAddr>,
59    ) -> Self {
60        let query = Self::create_query(&dht, topic, &opts, &keypair, &relay_addresses);
61        Self {
62            topic,
63            opts,
64            dht,
65            keypair,
66            state: PeerDiscoveryState::Querying(query),
67            first_round_complete: false,
68            relay_addresses,
69        }
70    }
71
72    pub fn topic(&self) -> IdBytes {
73        self.topic
74    }
75
76    pub fn is_first_round_complete(&self) -> bool {
77        self.first_round_complete
78    }
79
80    /// If relay addresses have changed, update them for subsequent announce queries.
81    pub fn maybe_set_relay_addresses(&mut self, relay_addresses: &[SocketAddr]) {
82        if self.opts.server() && relay_addresses != self.relay_addresses {
83            self.relay_addresses = relay_addresses.to_vec();
84        }
85    }
86
87    /// Restart the query immediately, aborting any sleep timer.
88    pub fn refresh(&mut self) {
89        let query = Self::create_query(
90            &self.dht,
91            self.topic,
92            &self.opts,
93            &self.keypair,
94            &self.relay_addresses,
95        );
96        self.state = PeerDiscoveryState::Querying(query);
97    }
98
99    fn create_query(
100        dht: &Dht,
101        topic: IdBytes,
102        opts: &JoinOpts,
103        keypair: &Keypair,
104        relay_addresses: &[SocketAddr],
105    ) -> Box<ActiveQuery> {
106        Box::new(if opts.server() {
107            ActiveQuery::Announce(dht.announce(topic, keypair.clone(), relay_addresses.to_vec()))
108        } else {
109            ActiveQuery::Lookup(
110                dht.lookup(topic, Commit::No)
111                    .expect("lookup creation failed"),
112            )
113        })
114    }
115
116    fn refresh_duration() -> Duration {
117        Duration::from_secs(REFRESH_BASE_SECS + rand::random::<u64>() % REFRESH_JITTER_SECS)
118    }
119}
120
121impl Stream for PeerDiscovery {
122    type Item = hyperdht::Result<Option<LookupResponse>>;
123
124    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
125        loop {
126            match &mut self.state {
127                PeerDiscoveryState::Querying(query) => match Pin::new(query.as_mut()).poll_next(cx)
128                {
129                    Poll::Ready(Some(result)) => {
130                        return Poll::Ready(Some(result));
131                    }
132                    Poll::Ready(None) => {
133                        self.first_round_complete = true;
134                        let duration = Self::refresh_duration();
135                        self.state =
136                            PeerDiscoveryState::Sleeping(Box::pin(tokio::time::sleep(duration)));
137                    }
138                    Poll::Pending => return Poll::Pending,
139                },
140                PeerDiscoveryState::Sleeping(timer) => match timer.as_mut().poll(cx) {
141                    Poll::Ready(()) => {
142                        let query = Self::create_query(
143                            &self.dht,
144                            self.topic,
145                            &self.opts,
146                            &self.keypair,
147                            &self.relay_addresses,
148                        );
149                        self.state = PeerDiscoveryState::Querying(query);
150                    }
151                    Poll::Pending => {
152                        return Poll::Pending;
153                    }
154                },
155            }
156        }
157    }
158}
hyperswarm/peer_discovery.rs

hyperswarm/
peer_discovery.rs
