reth_dns_discovery/
sync.rs

Help
1use crate::tree::{LinkEntry, TreeRootEntry};
2use enr::EnrKeyUnambiguous;
3use linked_hash_set::LinkedHashSet;
4use secp256k1::SecretKey;
5use std::{
6    collections::HashMap,
7    time::{Duration, Instant},
8};
9
10/// A sync-able tree
11pub(crate) struct SyncTree<K: EnrKeyUnambiguous = SecretKey> {
12    /// Root of the tree
13    root: TreeRootEntry,
14    /// Link to this tree
15    link: LinkEntry<K>,
16    /// Timestamp when the root was updated
17    root_updated: Instant,
18    /// The state of the tree sync progress.
19    sync_state: SyncState,
20    /// Links contained in this tree
21    resolved_links: HashMap<String, LinkEntry<K>>,
22    /// Unresolved links of the tree
23    unresolved_links: LinkedHashSet<String>,
24    /// Unresolved nodes of the tree
25    unresolved_nodes: LinkedHashSet<String>,
26}
27
28// === impl SyncTree ===
29
30impl<K: EnrKeyUnambiguous> SyncTree<K> {
31    pub(crate) fn new(root: TreeRootEntry, link: LinkEntry<K>) -> Self {
32        Self {
33            root,
34            link,
35            root_updated: Instant::now(),
36            sync_state: SyncState::Pending,
37            resolved_links: Default::default(),
38            unresolved_links: Default::default(),
39            unresolved_nodes: Default::default(),
40        }
41    }
42
43    #[cfg(test)]
44    pub(crate) const fn root(&self) -> &TreeRootEntry {
45        &self.root
46    }
47
48    pub(crate) const fn link(&self) -> &LinkEntry<K> {
49        &self.link
50    }
51
52    pub(crate) fn resolved_links_mut(&mut self) -> &mut HashMap<String, LinkEntry<K>> {
53        &mut self.resolved_links
54    }
55
56    pub(crate) fn extend_children(
57        &mut self,
58        kind: ResolveKind,
59        children: impl IntoIterator<Item = String>,
60    ) {
61        match kind {
62            ResolveKind::Enr => {
63                self.unresolved_nodes.extend(children);
64            }
65            ResolveKind::Link => {
66                self.unresolved_links.extend(children);
67            }
68        }
69    }
70
71    /// Advances the state of the tree by returning actions to perform
72    pub(crate) fn poll(&mut self, now: Instant, update_timeout: Duration) -> Option<SyncAction> {
73        match self.sync_state {
74            SyncState::Pending => {
75                self.sync_state = SyncState::Enr;
76                return Some(SyncAction::Link(self.root.link_root.clone()))
77            }
78            SyncState::Enr => {
79                self.sync_state = SyncState::Active;
80                return Some(SyncAction::Enr(self.root.enr_root.clone()))
81            }
82            SyncState::Link => {
83                self.sync_state = SyncState::Active;
84                return Some(SyncAction::Link(self.root.link_root.clone()))
85            }
86            SyncState::Active => {
87                if now > self.root_updated + update_timeout {
88                    self.sync_state = SyncState::RootUpdate;
89                    return Some(SyncAction::UpdateRoot)
90                }
91            }
92            SyncState::RootUpdate => return None,
93        }
94
95        if let Some(link) = self.unresolved_links.pop_front() {
96            return Some(SyncAction::Link(link))
97        }
98
99        let enr = self.unresolved_nodes.pop_front()?;
100        Some(SyncAction::Enr(enr))
101    }
102
103    /// Updates the root and returns what changed
104    pub(crate) fn update_root(&mut self, root: TreeRootEntry) {
105        let enr = root.enr_root == self.root.enr_root;
106        let link = root.link_root == self.root.link_root;
107
108        self.root = root;
109        self.root_updated = Instant::now();
110
111        let state = match (enr, link) {
112            (true, true) => {
113                self.unresolved_nodes.clear();
114                self.unresolved_links.clear();
115                SyncState::Pending
116            }
117            (true, _) => {
118                self.unresolved_nodes.clear();
119                SyncState::Enr
120            }
121            (_, true) => {
122                self.unresolved_links.clear();
123                SyncState::Link
124            }
125            _ => {
126                // unchanged
127                return
128            }
129        };
130        self.sync_state = state;
131    }
132}
133
134/// The action to perform by the service
135pub(crate) enum SyncAction {
136    UpdateRoot,
137    Enr(String),
138    Link(String),
139}
140
141/// How the [`SyncTree::update_root`] changed the root
142enum SyncState {
143    RootUpdate,
144    Pending,
145    Enr,
146    Link,
147    Active,
148}
149
150/// What kind of hash to resolve
151pub(crate) enum ResolveKind {
152    Enr,
153    Link,
154}
155
156// === impl ResolveKind ===
157
158impl ResolveKind {
159    pub(crate) const fn is_link(&self) -> bool {
160        matches!(self, Self::Link)
161    }
162}
reth_dns_discovery/sync.rs

reth_dns_discovery/
sync.rs
