Skip to content

Commit

Permalink
feat(iroh-net): Add StaticDiscovery to provide static info to endpoin…
Browse files Browse the repository at this point in the history
…ts (#2825)

## Description

Add StaticProvider to provide static info to endpoints.

The long term plan is to reduce ways to interact directly with the
endpoint node map and instead provide this.

## Breaking Changes

None, this just adds stuff. Long term we want to remove
endpoint.add_node_info with this.

## Notes & open questions

<!-- Any notes, remarks or open questions you have to make about the PR.
-->

## Change checklist

- [ ] Self-review.
- [ ] Documentation updates following the [style
guide](https://rust-lang.github.io/rfcs/1574-more-api-documentation-conventions.html#appendix-a-full-conventions-text),
if relevant.
- [ ] Tests if relevant.
- [ ] All breaking changes documented.
  • Loading branch information
rklaehn authored Oct 24, 2024
1 parent d29537d commit c9d1ba7
Show file tree
Hide file tree
Showing 2 changed files with 163 additions and 0 deletions.
1 change: 1 addition & 0 deletions iroh-net/src/discovery.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,7 @@ pub mod dns;
#[cfg_attr(iroh_docsrs, doc(cfg(feature = "discovery-local-network")))]
pub mod local_swarm_discovery;
pub mod pkarr;
pub mod static_provider;

/// Node discovery for [`super::Endpoint`].
///
Expand Down
162 changes: 162 additions & 0 deletions iroh-net/src/discovery/static_provider.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
//! A static discovery implementation that allows adding info for nodes manually.
use std::{
collections::{btree_map::Entry, BTreeMap},
sync::{Arc, RwLock},
time::SystemTime,
};

use futures_lite::stream::{self, StreamExt};
use iroh_base::{
key::NodeId,
node_addr::{AddrInfo, NodeAddr},
};

use super::{Discovery, DiscoveryItem};

/// A static discovery implementation that allows providing info for nodes manually.
#[derive(Debug, Default)]
#[repr(transparent)]
pub struct StaticProvider {
nodes: Arc<RwLock<BTreeMap<NodeId, NodeInfo>>>,
}

#[derive(Debug)]
struct NodeInfo {
info: AddrInfo,
last_updated: SystemTime,
}

impl StaticProvider {
/// The provenance string for this discovery implementation.
pub const PROVENANCE: &'static str = "static_discovery";

/// Create a new static discovery instance.
pub fn new() -> Self {
Self::default()
}

/// Creates a static discovery instance from something that can be converted into node addresses.
///
/// Example:
/// ```rust
/// use std::str::FromStr;
///
/// use iroh_base::ticket::NodeTicket;
/// use iroh_net::{Endpoint, discovery::static_provider::StaticProvider};
///
/// # async fn example() -> anyhow::Result<()> {
/// # #[derive(Default)] struct Args { tickets: Vec<NodeTicket> }
/// # let args = Args::default();
/// // get tickets from command line args
/// let tickets: Vec<NodeTicket> = args.tickets;
/// // create a StaticProvider from the tickets. Ticket info will be combined if multiple tickets refer to the same node.
/// let discovery = StaticProvider::from_node_addrs(tickets);
/// // create an endpoint with the discovery
/// let endpoint = Endpoint::builder()
/// .add_discovery(|_| Some(discovery))
/// .bind().await?;
/// # Ok(())
/// # }
/// ```
pub fn from_node_addrs(infos: impl IntoIterator<Item = impl Into<NodeAddr>>) -> Self {
let res = Self::default();
for info in infos {
res.add_node_addr(info);
}
res
}

/// Add node info for the given node id.
///
/// This will completely overwrite any existing info for the node.
pub fn set_node_addr(&self, info: impl Into<NodeAddr>) -> Option<NodeAddr> {
let last_updated = SystemTime::now();
let info: NodeAddr = info.into();
let mut guard = self.nodes.write().unwrap();
let previous = guard.insert(
info.node_id,
NodeInfo {
info: info.info,
last_updated,
},
);
previous.map(|x| NodeAddr {
node_id: info.node_id,
info: x.info,
})
}

/// Add node info for the given node id, combining it with any existing info.
///
/// This will add any new direct addresses and overwrite the relay url.
pub fn add_node_addr(&self, info: impl Into<NodeAddr>) {
let info: NodeAddr = info.into();
let last_updated = SystemTime::now();
let mut guard = self.nodes.write().unwrap();
match guard.entry(info.node_id) {
Entry::Occupied(mut entry) => {
let existing = entry.get_mut();
existing
.info
.direct_addresses
.extend(info.info.direct_addresses);
existing.info.relay_url = info.info.relay_url;
existing.last_updated = last_updated;
}
Entry::Vacant(entry) => {
entry.insert(NodeInfo {
info: info.info,
last_updated,
});
}
}
}

/// Get node info for the given node id.
pub fn get_node_addr(&self, node_id: NodeId) -> Option<NodeAddr> {
let guard = self.nodes.read().unwrap();
let info = guard.get(&node_id).map(|x| x.info.clone())?;
Some(NodeAddr { node_id, info })
}

/// Remove node info for the given node id.
pub fn remove_node_addr(&self, node_id: NodeId) -> Option<NodeAddr> {
let mut guard = self.nodes.write().unwrap();
let res = guard.remove(&node_id)?;
Some(NodeAddr {
node_id,
info: res.info,
})
}
}

impl Discovery for StaticProvider {
fn publish(&self, _info: &AddrInfo) {}

fn resolve(
&self,
_endpoint: crate::Endpoint,
node_id: NodeId,
) -> Option<futures_lite::stream::Boxed<anyhow::Result<super::DiscoveryItem>>> {
let guard = self.nodes.read().unwrap();
let info = guard.get(&node_id);
match info {
Some(addr_info) => {
let item = DiscoveryItem {
node_id,
provenance: Self::PROVENANCE,
last_updated: Some(
addr_info
.last_updated
.duration_since(SystemTime::UNIX_EPOCH)
.expect("time drift")
.as_micros() as u64,
),
addr_info: addr_info.info.clone(),
};
Some(stream::iter(Some(Ok(item))).boxed())
}
None => None,
}
}
}

0 comments on commit c9d1ba7

Please sign in to comment.