feat: add DiscoveryItem::user_data method and adjust locally-discovered-nodes example

Author: “ramfox”

iroh/examples/locally-discovered-nodes.rs

@@ -5,66 +5,77 @@
 //! This is an async, non-determinate process, so the number of NodeIDs discovered each time may be different. If you have other iroh endpoints or iroh nodes with [`LocalSwarmDiscovery`] enabled, it may discover those nodes as well.
 use std::time::Duration;
 
-use iroh::{
-    discovery::local_swarm_discovery::LocalSwarmDiscovery, endpoint::Source, Endpoint, SecretKey,
-};
+use anyhow::Result;
+use iroh::{node_info::UserData, Endpoint, NodeId};
+use n0_future::StreamExt;
+use tokio::task::JoinSet;
 
 #[tokio::main]
-async fn main() -> anyhow::Result<()> {
+async fn main() -> Result<()> {
     tracing_subscriber::fmt::init();
-    println!("locally discovered nodes example!\n");
-    let mut rng = rand::rngs::OsRng;
-    let key = SecretKey::generate(&mut rng);
-    let id = key.public();
-    println!("creating endpoint {id:?}\n");
-    let ep = Endpoint::builder()
-        .secret_key(key)
-        .discovery(Box::new(LocalSwarmDiscovery::new(id)?))
-        .bind()
-        .await?;
+    println!("Discovering Local Nodes Example!");
 
-    let node_count = 5;
-    println!("creating {node_count} additional endpoints to discover locally:");
-    let mut discoverable_eps = Vec::with_capacity(node_count);
-    for _ in 0..node_count {
-        let key = SecretKey::generate(&mut rng);
-        let id = key.public();
-        println!("\t{id:?}");
-        let ep = Endpoint::builder()
-            .secret_key(key)
-            .discovery(Box::new(LocalSwarmDiscovery::new(id)?))
-            .bind()
-            .await?;
-        discoverable_eps.push(ep);
-    }
+    let ep = Endpoint::builder().discovery_local_network().bind().await?;
+    let node_id = ep.node_id();
+    println!("Created endpoint {}", node_id.fmt_short());
 
-    let duration = Duration::from_secs(3);
-    println!("\nwaiting {duration:?} to allow discovery to occur...\n");
-    tokio::time::sleep(duration).await;
+    let user_data = UserData::try_from(String::from("local-nodes-example"))?;
 
-    // get an iterator of all the remote nodes this endpoint knows about
-    let remotes = ep.remote_info_iter();
-    // filter that list down to the nodes that have a `Source::Discovery` with
-    // the `service` name [`iroh::discovery::local_swarm_discovery::NAME`]
-    // If you have a long running node and want to only get the nodes that were
-    // discovered recently, you can also filter on the `Duration` of the source,
-    // which indicates how long ago we got information from that source.
-    let locally_discovered: Vec<_> = remotes
-        .filter(|remote| {
-            remote.sources().iter().any(|(source, _duration)| {
-                if let Source::Discovery { name } = source {
-                    name == iroh::discovery::local_swarm_discovery::NAME
-                } else {
-                    false
+    let mut discovery_stream = ep.discovery_stream();
+
+    let ud = user_data.clone();
+    let discovery_stream_task = tokio::spawn(async move {
+        let mut discovered_nodes: Vec<NodeId> = vec![];
+        while let Some(item) = discovery_stream.next().await {
+            match item {
+                Err(e) => {
+                    tracing::error!("{e}");
+                    return;
                 }
-            })
-        })
-        .map(|remote| remote.node_id)
-        .collect();
+                Ok(item) => {
+                    // if there is no user data, or the user data
+                    // does not indicate that the discovered node
+                    // is a part of the example, ignore it
+                    match item.node_info().data.user_data() {
+                        Some(user_data) if &ud == user_data => {}
+                        _ => {
+                            tracing::error!("found node with unexpected user data, ignoring it");
+                            continue;
+                        }
+                    }
 
-    println!("found:");
-    for id in locally_discovered {
-        println!("\t{id:?}");
+                    // if we've already found this node, ignore it
+                    // otherwise announce that we have found a new node
+                    if discovered_nodes.contains(&item.node_id()) {
+                        continue;
+                    } else {
+                        discovered_nodes.push(item.node_id());
+                        println!("Found node {}!", item.node_id().fmt_short());
+                    }
+                }
+            };
+        }
+    });
+
+    let mut set = JoinSet::new();
+    let node_count = 5;
+    for _ in 0..node_count {
+        let ud = user_data.clone();
+        set.spawn(async move {
+            let ep = Endpoint::builder().discovery_local_network().bind().await?;
+            ep.set_user_data_for_discovery(Some(ud));
+            tokio::time::sleep(Duration::from_secs(3)).await;
+            ep.close().await;
+            anyhow::Ok(())
+        });
     }
+
+    set.join_all().await.iter().for_each(|res| {
+        if let Err(e) = res {
+            tracing::error!("{e}");
+        }
+    });
+    ep.close().await;
+    discovery_stream_task.abort();
     Ok(())
 }

iroh/src/discovery.rs

@@ -264,6 +264,11 @@ impl DiscoveryItem {
     pub fn into_node_addr(self) -> NodeAddr {
         self.node_info.into_node_addr()
     }
+
+    /// Returns any user-defined data.
+    pub fn user_data(&self) -> Option<UserData> {
+        self.node_info().data.user_data().cloned()
+    }
 }
 
 impl std::ops::Deref for DiscoveryItem {