Skip to main content

reth_dns_discovery/
sync.rs

1use crate::tree::{LinkEntry, TreeRootEntry};
2use enr::EnrKeyUnambiguous;
3use linked_hash_set::LinkedHashSet;
4use secp256k1::SecretKey;
5use std::time::{Duration, Instant};
6
7/// A sync-able tree
8pub(crate) struct SyncTree<K: EnrKeyUnambiguous = SecretKey> {
9    /// Root of the tree
10    root: TreeRootEntry,
11    /// Link to this tree
12    link: LinkEntry<K>,
13    /// Timestamp when the root was updated
14    root_updated: Instant,
15    /// The state of the tree sync progress.
16    sync_state: SyncState,
17    /// Unresolved links of the tree
18    unresolved_links: LinkedHashSet<String>,
19    /// Unresolved nodes of the tree
20    unresolved_nodes: LinkedHashSet<String>,
21}
22
23// === impl SyncTree ===
24
25impl<K: EnrKeyUnambiguous> SyncTree<K> {
26    pub(crate) fn new(root: TreeRootEntry, link: LinkEntry<K>) -> Self {
27        Self {
28            root,
29            link,
30            root_updated: Instant::now(),
31            sync_state: SyncState::Pending,
32            unresolved_links: Default::default(),
33            unresolved_nodes: Default::default(),
34        }
35    }
36
37    #[cfg(test)]
38    pub(crate) const fn root(&self) -> &TreeRootEntry {
39        &self.root
40    }
41
42    pub(crate) const fn link(&self) -> &LinkEntry<K> {
43        &self.link
44    }
45
46    pub(crate) fn extend_children(
47        &mut self,
48        kind: ResolveKind,
49        children: impl IntoIterator<Item = String>,
50    ) {
51        match kind {
52            ResolveKind::Enr => {
53                self.unresolved_nodes.extend(children);
54            }
55            ResolveKind::Link => {
56                self.unresolved_links.extend(children);
57            }
58        }
59    }
60
61    /// Advances the state of the tree by returning actions to perform
62    pub(crate) fn poll(&mut self, now: Instant, update_timeout: Duration) -> Option<SyncAction> {
63        match self.sync_state {
64            SyncState::Pending => {
65                self.sync_state = SyncState::Enr;
66                return Some(SyncAction::Link(self.root.link_root.clone()))
67            }
68            SyncState::Enr => {
69                self.sync_state = SyncState::Active;
70                return Some(SyncAction::Enr(self.root.enr_root.clone()))
71            }
72            SyncState::Link => {
73                self.sync_state = SyncState::Active;
74                return Some(SyncAction::Link(self.root.link_root.clone()))
75            }
76            SyncState::Active => {
77                if now > self.root_updated + update_timeout {
78                    self.sync_state = SyncState::RootUpdate;
79                    return Some(SyncAction::UpdateRoot)
80                }
81            }
82            SyncState::RootUpdate => return None,
83        }
84
85        if let Some(link) = self.unresolved_links.pop_front() {
86            return Some(SyncAction::Link(link))
87        }
88
89        let enr = self.unresolved_nodes.pop_front()?;
90        Some(SyncAction::Enr(enr))
91    }
92
93    /// Records that the scheduled root lookup failed.
94    ///
95    /// Without this the tree would stay in [`SyncState::RootUpdate`] and stop polling entirely,
96    /// so one failed lookup would take the tree out of rotation permanently. It retries at the
97    /// next recheck instead.
98    pub(crate) fn root_update_failed(&mut self) {
99        // Only a lookup this tree scheduled releases it. Root lookups are also started from
100        // bootstrap, from the public sync command and from link entries found in other trees, and
101        // those outcomes must not overwrite a walk that is still recorded in `Pending`/`Enr`/
102        // `Link`, or the tree would go quiet with its subtree unvisited.
103        if matches!(self.sync_state, SyncState::RootUpdate) {
104            self.root_updated = Instant::now();
105            self.sync_state = SyncState::Active;
106        }
107    }
108
109    /// Updates the root and returns what changed
110    pub(crate) fn update_root(&mut self, root: TreeRootEntry) {
111        let enr_unchanged = root.enr_root == self.root.enr_root;
112        let link_unchanged = root.link_root == self.root.link_root;
113
114        self.root = root;
115        self.root_updated = Instant::now();
116
117        let state = match (enr_unchanged, link_unchanged) {
118            // Nothing to resync. The tree still has to leave `RootUpdate`, which would otherwise
119            // park `poll` on `None` forever, but only when this answers its own recheck: an
120            // unsolicited lookup of an unchanged root must leave a walk in progress alone.
121            (true, true) => {
122                if !matches!(self.sync_state, SyncState::RootUpdate) {
123                    return
124                }
125                SyncState::Active
126            }
127            // only ENR changed
128            (false, true) => {
129                self.unresolved_nodes.clear();
130                SyncState::Enr
131            }
132            // only LINK changed
133            (true, false) => {
134                self.unresolved_links.clear();
135                SyncState::Link
136            }
137            // both changed
138            (false, false) => {
139                self.unresolved_nodes.clear();
140                self.unresolved_links.clear();
141                SyncState::Pending
142            }
143        };
144        self.sync_state = state;
145    }
146}
147
148/// The action to perform by the service
149#[derive(Debug)]
150pub(crate) enum SyncAction {
151    UpdateRoot,
152    Enr(String),
153    Link(String),
154}
155
156/// How the [`SyncTree::update_root`] changed the root
157enum SyncState {
158    RootUpdate,
159    Pending,
160    Enr,
161    Link,
162    Active,
163}
164
165/// What kind of hash to resolve
166pub(crate) enum ResolveKind {
167    Enr,
168    Link,
169}
170
171// === impl ResolveKind ===
172
173impl ResolveKind {
174    pub(crate) const fn is_link(&self) -> bool {
175        matches!(self, Self::Link)
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182    use enr::EnrKey;
183    use secp256k1::rand::thread_rng;
184
185    fn base_root() -> TreeRootEntry {
186        // taken from existing tests to ensure valid formatting
187        let s = "enrtree-root:v1 e=QFT4PBCRX4XQCV3VUYJ6BTCEPU l=JGUFMSAGI7KZYB3P7IZW4S5Y3A seq=3 sig=3FmXuVwpa8Y7OstZTx9PIb1mt8FrW7VpDOFv4AaGCsZ2EIHmhraWhe4NxYhQDlw5MjeFXYMbJjsPeKlHzmJREQE";
188        s.parse::<TreeRootEntry>().unwrap()
189    }
190
191    fn make_tree() -> SyncTree {
192        let secret_key = SecretKey::new(&mut thread_rng());
193        let link =
194            LinkEntry { domain: "nodes.example.org".to_string(), pubkey: secret_key.public() };
195        SyncTree::new(base_root(), link)
196    }
197
198    fn advance_to_active(tree: &mut SyncTree) {
199        // Move Pending -> (emit Link) -> Enr, then Enr -> (emit Enr) -> Active
200        let now = Instant::now();
201        let timeout = Duration::from_secs(60 * 60 * 24);
202        let _ = tree.poll(now, timeout);
203        let _ = tree.poll(now, timeout);
204    }
205
206    #[test]
207    fn update_root_unchanged_no_action_from_active() {
208        let mut tree = make_tree();
209        let now = Instant::now();
210        let timeout = Duration::from_secs(60 * 60 * 24);
211        advance_to_active(&mut tree);
212
213        // same root -> no resync
214        let same = base_root();
215        tree.update_root(same);
216        assert!(tree.poll(now, timeout).is_none());
217    }
218
219    #[test]
220    fn unsolicited_root_outcome_does_not_cancel_a_walk() {
221        let timeout = Duration::from_secs(60);
222
223        // a root lookup this tree did not schedule must not discard the walk it still owes. Root
224        // lookups are also started from bootstrap, the public sync command and link entries in
225        // other trees, and the query pool does not de-duplicate them.
226        for unsolicited in [true, false] {
227            let mut tree = make_tree();
228            // Pending -> emits Link, leaving the enr root still to walk
229            assert!(matches!(tree.poll(Instant::now(), timeout), Some(SyncAction::Link(_))));
230
231            if unsolicited {
232                tree.update_root(base_root());
233            } else {
234                tree.root_update_failed();
235            }
236
237            assert!(
238                matches!(tree.poll(Instant::now(), timeout), Some(SyncAction::Enr(_))),
239                "walk was cancelled by an unsolicited root outcome"
240            );
241        }
242    }
243
244    #[test]
245    fn update_root_unchanged_still_rechecks_later() {
246        let mut tree = make_tree();
247        let timeout = Duration::from_secs(60);
248        advance_to_active(&mut tree);
249
250        // the recheck interval elapses and the root is looked up again
251        let due = Instant::now() + timeout * 2;
252        assert!(matches!(tree.poll(due, timeout), Some(SyncAction::UpdateRoot)));
253
254        // it comes back unchanged, so there is nothing to resync right now
255        tree.update_root(base_root());
256        assert!(tree.poll(Instant::now(), timeout).is_none());
257
258        // but the tree has to keep rechecking, otherwise it never picks up later changes
259        let later = Instant::now() + timeout * 4;
260        assert!(
261            matches!(tree.poll(later, timeout), Some(SyncAction::UpdateRoot)),
262            "tree stopped rechecking its root after an unchanged update"
263        );
264    }
265
266    #[test]
267    fn failed_root_update_still_rechecks_later() {
268        let mut tree = make_tree();
269        let timeout = Duration::from_secs(60);
270        advance_to_active(&mut tree);
271
272        let due = Instant::now() + timeout * 2;
273        assert!(matches!(tree.poll(due, timeout), Some(SyncAction::UpdateRoot)));
274
275        // the lookup failed, which must not take the tree out of rotation for good
276        tree.root_update_failed();
277        assert!(tree.poll(Instant::now(), timeout).is_none());
278
279        let later = Instant::now() + timeout * 4;
280        assert!(
281            matches!(tree.poll(later, timeout), Some(SyncAction::UpdateRoot)),
282            "tree stopped rechecking its root after a failed lookup"
283        );
284    }
285
286    #[test]
287    fn update_root_only_enr_changed_triggers_enr() {
288        let mut tree = make_tree();
289        advance_to_active(&mut tree);
290        let mut new_root = base_root();
291        new_root.enr_root = "NEW_ENR_ROOT".to_string();
292        let now = Instant::now();
293        let timeout = Duration::from_secs(60 * 60 * 24);
294
295        tree.update_root(new_root.clone());
296        match tree.poll(now, timeout) {
297            Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
298            other => panic!("expected Enr action, got {:?}", other),
299        }
300    }
301
302    #[test]
303    fn update_root_only_link_changed_triggers_link() {
304        let mut tree = make_tree();
305        advance_to_active(&mut tree);
306        let mut new_root = base_root();
307        new_root.link_root = "NEW_LINK_ROOT".to_string();
308        let now = Instant::now();
309        let timeout = Duration::from_secs(60 * 60 * 24);
310
311        tree.update_root(new_root.clone());
312        match tree.poll(now, timeout) {
313            Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
314            other => panic!("expected Link action, got {:?}", other),
315        }
316    }
317
318    #[test]
319    fn update_root_both_changed_triggers_link_then_enr() {
320        let mut tree = make_tree();
321        advance_to_active(&mut tree);
322        let mut new_root = base_root();
323        new_root.enr_root = "NEW_ENR_ROOT".to_string();
324        new_root.link_root = "NEW_LINK_ROOT".to_string();
325        let now = Instant::now();
326        let timeout = Duration::from_secs(60 * 60 * 24);
327
328        tree.update_root(new_root.clone());
329        match tree.poll(now, timeout) {
330            Some(SyncAction::Link(hash)) => assert_eq!(hash, new_root.link_root),
331            other => panic!("expected first Link action, got {:?}", other),
332        }
333        match tree.poll(now, timeout) {
334            Some(SyncAction::Enr(hash)) => assert_eq!(hash, new_root.enr_root),
335            other => panic!("expected second Enr action, got {:?}", other),
336        }
337    }
338}