Merge 11280b5 into 14795ab

Author: ramfox

iroh/examples/locally-discovered-nodes.rs

@@ -1,70 +1,81 @@
-//! A small example showing how to get a list of nodes that were discovered via [`iroh::discovery::LocalSwarmDiscovery`]. LocalSwarmDiscovery uses [`swarm-discovery`](https://crates.io/crates/swarm-discovery) to discover other nodes in the local network ala mDNS.
+//! A small example showing how to get a list of nodes that were discovered via [`iroh::discovery::MdnsDiscovery`]. MdnsDiscovery uses [`swarm-discovery`](https://crates.io/crates/swarm-discovery), an opinionated implementation of mDNS to discover other nodes in the local network.
 //!
 //! This example creates an iroh endpoint, a few additional iroh endpoints to discover, waits a few seconds, and reports all of the iroh NodeIds (also called `[iroh::key::PublicKey]`s) it has discovered.
 //!
-//! 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.
+//! 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 [`MdnsDiscovery`] 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

@@ -34,8 +34,8 @@
 //! - The [`PkarrResolver`] which can perform lookups from designated [pkarr relay servers]
 //!   using HTTP.
 //!
-//! - [`LocalSwarmDiscovery`]: local_swarm_discovery::LocalSwarmDiscovery which is an mDNS
-//!   implementation.
+//! - [`MdnsDiscovery`]: mdns::MdnsDiscovery which uses the crate `swarm-discovery`, an
+//!   opinionated mDNS implementation, to discover nodes on the local network.
 //!
 //! - The [`DhtDiscovery`] also uses the [`pkarr`] system but can also publish and lookup
 //!   records to/from the Mainline DHT.
@@ -69,14 +69,14 @@
 //! # }
 //! ```
 //!
-//! To also enable [`LocalSwarmDiscovery`] it can be added as another service in the
+//! To also enable [`MdnsDiscovery`] it can be added as another service in the
 //! [`ConcurrentDiscovery`]:
 //!
 //! ```no_run
 //! # #[cfg(feature = "discovery-local-network")]
 //! # {
 //! # use iroh::discovery::dns::DnsDiscovery;
-//! # use iroh::discovery::local_swarm_discovery::LocalSwarmDiscovery;
+//! # use iroh::discovery::mdns::MdnsDiscovery;
 //! # use iroh::discovery::pkarr::PkarrPublisher;
 //! # use iroh::discovery::ConcurrentDiscovery;
 //! # use iroh::SecretKey;
@@ -86,7 +86,7 @@
 //! let discovery = ConcurrentDiscovery::from_services(vec![
 //!     Box::new(PkarrPublisher::n0_dns(secret_key.clone())),
 //!     Box::new(DnsDiscovery::n0_dns()),
-//!     Box::new(LocalSwarmDiscovery::new(secret_key.public())?),
+//!     Box::new(MdnsDiscovery::new(secret_key.public())?),
 //! ]);
 //! # Ok(())
 //! # }
@@ -102,7 +102,7 @@
 //! [`PkarrPublisher`]: pkarr::PkarrPublisher
 //! [`DhtDiscovery`]: pkarr::dht::DhtDiscovery
 //! [pkarr relay servers]: https://pkarr.org/#servers
-//! [`LocalSwarmDiscovery`]: local_swarm_discovery::LocalSwarmDiscovery
+//! [`MdnsDiscovery`]: mdns::MdnsDiscovery
 //! [`StaticProvider`]: static_provider::StaticProvider
 
 use std::sync::Arc;
@@ -126,7 +126,7 @@ use crate::Endpoint;
 pub mod dns;
 
 #[cfg(feature = "discovery-local-network")]
-pub mod local_swarm_discovery;
+pub mod mdns;
 pub mod pkarr;
 pub mod static_provider;
 
@@ -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 {