Преглед изворни кода

dht: DhtHandler::node() returns a Result<Node> instead of Node

epiphany пре 8 месеци
родитељ
комит
05ae3878e1
3 измењених фајлова са 34 додато и 15 уклоњено
  1. 1 1
      src/dht/handler.rs
  2. 29 14
      src/dht/mod.rs
  3. 4 0
      src/dht/tasks.rs

+ 1 - 1
src/dht/handler.rs

@@ -37,7 +37,7 @@ pub trait DhtHandler: Send + Sync + Sized {
     fn dht(&self) -> Arc<Dht<Self>>;
 
     /// Get our own node
-    async fn node(&self) -> Self::Node;
+    async fn node(&self) -> Result<Self::Node>;
 
     /// Send PING request, which is used to know the node data of a peer
     /// (and most importantly, its ID/key in the DHT keyspace)

+ 29 - 14
src/dht/mod.rs

@@ -305,7 +305,7 @@ impl<H: DhtHandler> Dht<H> {
         value: &H::Value,
         message: &M,
     ) -> Result<()> {
-        let self_node = self.handler().await.node().await;
+        let self_node = self.handler().await.node().await?;
         if self_node.addresses().is_empty() {
             return Err(().into()); // TODO
         }
@@ -326,19 +326,25 @@ impl<H: DhtHandler> Dht<H> {
 
     /// Lookup our own node id
     pub async fn bootstrap(&self) {
+        let self_node = self.handler().await.node().await;
+        if self_node.is_err() {
+            return;
+        }
+        let self_node = self_node.unwrap();
+
         self.set_bootstrapped(true).await;
 
         info!(target: "dht::bootstrap()", "[DHT] Bootstrapping");
         self.event_publisher.notify(DhtEvent::BootstrapStarted).await;
 
-        let self_node_id = self.handler().await.node().await.id();
-        let nodes = self.lookup_nodes(&self_node_id).await;
+        let _nodes = self.lookup_nodes(&self_node.id()).await;
 
-        if nodes.is_empty() {
-            self.set_bootstrapped(false).await;
-        } else {
-            self.event_publisher.notify(DhtEvent::BootstrapCompleted).await;
-        }
+        // if nodes.is_empty() {
+        //     self.set_bootstrapped(false).await;
+        // } else {
+        // }
+
+        self.event_publisher.notify(DhtEvent::BootstrapCompleted).await;
     }
 
     // TODO: Optimize this
@@ -351,7 +357,11 @@ impl<H: DhtHandler> Dht<H> {
         }
 
         // Send keys that are closer to this node than we are
-        let self_id = self.handler().await.node().await.id();
+        let self_node = self.handler().await.node().await;
+        if self_node.is_err() {
+            return;
+        }
+        let self_id = self_node.unwrap().id();
         for (key, value) in self.hash_table.read().await.iter() {
             let node_distance = BigUint::from_bytes_be(&self.distance(key, &node.id()));
             let self_distance = BigUint::from_bytes_be(&self.distance(key, &self_id));
@@ -375,7 +385,10 @@ impl<H: DhtHandler> Dht<H> {
     pub async fn remove_node(&self, node_id: &blake3::Hash) {
         let handler = self.handler().await;
         let self_node = handler.node().await;
-        let bucket_index = handler.dht().get_bucket_index(&self_node.id(), node_id).await;
+        if self_node.is_err() {
+            return;
+        }
+        let bucket_index = handler.dht().get_bucket_index(&self_node.unwrap().id(), node_id).await;
         let buckets_lock = handler.dht().buckets.clone();
         let mut buckets = buckets_lock.write().await;
         let bucket = &mut buckets[bucket_index];
@@ -550,10 +563,12 @@ impl<H: DhtHandler> Dht<H> {
                         .await;
 
                     // Remove our own node and duplicates
-                    let self_id = self.handler().await.node().await.id();
-                    nodes.retain(|node: &H::Node| {
-                        node.id() != self_id && seen_nodes.insert(node.id())
-                    });
+                    if let Ok(self_node) = self.handler().await.node().await {
+                        let self_id = self_node.id();
+                        nodes.retain(|node: &H::Node| {
+                            node.id() != self_id && seen_nodes.insert(node.id())
+                        });
+                    }
 
                     // Add new nodes to the list of nodes to visit
                     nodes_to_visit.extend(nodes.clone());

+ 4 - 0
src/dht/tasks.rs

@@ -182,6 +182,10 @@ pub async fn add_node_task<H: DhtHandler>(handler: Arc<H>) -> Result<()> {
         let (node, channel) = dht.add_node_rx.recv().await.unwrap();
 
         let self_node = handler.node().await;
+        if self_node.is_err() {
+            continue;
+        }
+        let self_node = self_node.unwrap();
 
         let bucket_index = dht.get_bucket_index(&self_node.id(), &node.id()).await;
         let buckets_lock = dht.buckets.clone();