use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, AtomicI32, AtomicUsize, Ordering};
use std::sync::Arc;
use futures_util::StreamExt;
use iroh::endpoint::VarInt;
use iroh::{Endpoint, EndpointAddr, EndpointId, RelayMode, SecretKey, TransportAddr};
use iroh_gossip::api::{Event, GossipReceiver, GossipSender, JoinOptions};
use iroh_gossip::net::{Gossip, GOSSIP_ALPN};
use iroh_gossip::proto::TopicId;
use tokio::sync::{mpsc, watch};
use super::wire::{self, NODE_KEY_LEN};
#[derive(Debug, Clone)]
pub enum MeshEvent {
Data { payload: Vec<u8>, sender: [u8; NODE_KEY_LEN], direct: bool },
NeighborUp([u8; NODE_KEY_LEN]),
NeighborDown([u8; NODE_KEY_LEN]),
Lagged,
}
pub struct Mesh {
endpoint: Endpoint,
gossip: Gossip,
key: [u8; NODE_KEY_LEN],
topics: tokio::sync::RwLock<HashMap<TopicId, Topic>>,
closing: Arc<std::sync::Mutex<HashMap<TopicId, watch::Receiver<bool>>>>,
leases: AtomicUsize,
}
pub struct MeshLease(Arc<Mesh>);
impl MeshLease {
pub fn mesh(&self) -> &Arc<Mesh> {
&self.0
}
}
impl Drop for MeshLease {
fn drop(&mut self) {
self.0.leases.fetch_sub(1, Ordering::SeqCst);
}
}
struct Topic {
sender: GossipSender,
pump: crate::rt::JoinHandle<()>,
seq: AtomicI32,
neighbors: Arc<AtomicUsize>,
}
const ADD_PEER_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
fn alpns() -> Vec<Vec<u8>> {
#[allow(unused_mut)]
let mut alpns = vec![GOSSIP_ALPN.to_vec()];
#[cfg(feature = "calls")]
alpns.push(crate::calls::wire::CALL_ALPN.to_vec());
alpns
}
static MESH: tokio::sync::Mutex<Option<(u64, Arc<Mesh>)>> = tokio::sync::Mutex::const_new(None);
pub async fn mesh() -> Result<Arc<Mesh>, String> {
acquire(false).await
}
pub async fn lease() -> Result<MeshLease, String> {
acquire(true).await.map(MeshLease)
}
async fn acquire(lease: bool) -> Result<Arc<Mesh>, String> {
let session = crate::db::current_session();
let mut slot = MESH.lock().await;
match claim(&mut slot, session.id(), session.is_live()) {
Claim::Reuse(m) => {
if lease {
m.leases.fetch_add(1, Ordering::SeqCst);
}
return Ok(m);
}
Claim::Refuse => return Err("the account is switching".into()),
Claim::Bind(previous) => {
if let Some(old) = previous {
old.close_detached();
}
}
}
let m = Arc::new(Mesh::bind().await.map_err(|e| format!("Realtime unavailable: {e}"))?);
if lease {
m.leases.fetch_add(1, Ordering::SeqCst);
}
*slot = Some((session.id(), m.clone()));
Ok(m)
}
pub async fn retire_if_idle(expected: &Arc<Mesh>) -> bool {
let mut slot = MESH.lock().await;
let idle = match slot.as_ref() {
Some((_, m)) if Arc::ptr_eq(m, expected) => {
m.leases.load(Ordering::SeqCst) == 0 && m.topics.read().await.is_empty() && !call_active()
}
_ => false,
};
if !idle {
return false;
}
if let Some((_, m)) = slot.take() {
m.close_detached();
}
true
}
#[cfg(feature = "calls")]
fn call_active() -> bool {
crate::calls::session::snapshot().is_some()
}
#[cfg(not(feature = "calls"))]
fn call_active() -> bool {
false
}
pub fn retire(owner: u64) {
if !crate::rt::can_spawn() {
return;
}
crate::rt::spawn(async move {
if let Some(m) = take_owned(&mut *MESH.lock().await, owner) {
m.close_detached();
}
});
}
#[cfg(test)]
pub(crate) async fn install_for_test(owner: u64, mesh: Mesh) -> Arc<Mesh> {
let mesh = Arc::new(mesh);
*MESH.lock().await = Some((owner, mesh.clone()));
mesh
}
#[cfg(test)]
pub(crate) static SLOT_TESTS: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
#[cfg(test)]
pub(crate) async fn slot_owner() -> Option<u64> {
MESH.lock().await.as_ref().map(|(owner, _)| *owner)
}
enum Claim<T> {
Reuse(T),
Bind(Option<T>),
Refuse,
}
fn claim<T: Clone>(slot: &mut Option<(u64, T)>, owner: u64, live: bool) -> Claim<T> {
match slot.as_ref() {
Some((o, m)) if *o == owner => Claim::Reuse(m.clone()),
_ if !live => Claim::Refuse,
_ => Claim::Bind(slot.take().map(|(_, m)| m)),
}
}
fn take_owned<T>(slot: &mut Option<(u64, T)>, owner: u64) -> Option<T> {
if slot.as_ref().is_some_and(|(o, _)| *o == owner) {
slot.take().map(|(_, m)| m)
} else {
None
}
}
impl Mesh {
async fn bind() -> anyhow_lite::Result<Self> {
let mesh = Self::build(wire::relay_mode()).await?;
let _ = crate::rt::time::timeout(std::time::Duration::from_secs(8), mesh.endpoint.online()).await;
for _ in 0..20 {
if mesh.endpoint.addr().addrs.iter().any(|a| matches!(a, TransportAddr::Relay(_))) {
break;
}
crate::rt::time::sleep(std::time::Duration::from_millis(100)).await;
}
crate::log_info!("[xdc] node {} via {}", short(&mesh.key), wire::relay_urls(&mesh.endpoint.addr()));
Ok(mesh)
}
#[cfg(test)]
pub(crate) async fn offline() -> Self {
Self::build(RelayMode::Disabled).await.expect("bind")
}
async fn build(relay_mode: RelayMode) -> anyhow_lite::Result<Self> {
let mut seed = [0u8; 32];
rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut seed);
let secret_key = SecretKey::from(seed);
let key = *secret_key.public().as_bytes();
let transport = iroh::endpoint::QuicTransportConfig::builder()
.keep_alive_interval(std::time::Duration::from_secs(15))
.max_idle_timeout(Some(std::time::Duration::from_secs(120).try_into()?))
.stream_receive_window(VarInt::from_u32(512 * 1024))
.receive_window(VarInt::from_u32(2 * 1024 * 1024))
.send_window(1_572_864)
.max_concurrent_bidi_streams(VarInt::from_u32(256))
.max_concurrent_uni_streams(VarInt::from_u32(256))
.initial_rtt(std::time::Duration::from_millis(100))
.congestion_controller_factory(Arc::new(noq_proto::congestion::CubicConfig::default()))
.send_observed_address_reports(false)
.receive_observed_address_reports(false)
.build();
struct RelayOnly(RelayMode);
impl iroh::endpoint::presets::Preset for RelayOnly {
fn apply(self, builder: iroh::endpoint::Builder) -> iroh::endpoint::Builder {
builder.relay_mode(self.0).crypto_provider(Arc::new(rustls::crypto::ring::default_provider()))
}
}
let endpoint = Endpoint::builder(RelayOnly(relay_mode))
.secret_key(secret_key)
.alpns(alpns())
.transport_config(transport)
.bind()
.await?;
let gossip = Gossip::builder().max_message_size(wire::MAX_GOSSIP_MESSAGE).spawn(endpoint.clone());
let (ep, g) = (endpoint.clone(), gossip.clone());
crate::rt::spawn(async move {
while let Some(incoming) = ep.accept().await {
let g = g.clone();
crate::rt::spawn(async move {
match incoming.await {
Ok(conn) if conn.alpn() == GOSSIP_ALPN => {
if let Err(e) = g.handle_connection(conn).await {
crate::log_debug!("[xdc] inbound gossip connection ended: {e}");
}
}
#[cfg(feature = "calls")]
Ok(conn) if conn.alpn() == crate::calls::wire::CALL_ALPN => {
crate::calls::session::on_incoming(conn).await;
}
Ok(_) => {}
Err(e) => crate::log_debug!("[xdc] accept failed: {e}"),
}
});
}
});
Ok(Self {
endpoint,
gossip,
key,
topics: tokio::sync::RwLock::new(HashMap::new()),
closing: Arc::new(std::sync::Mutex::new(HashMap::new())),
leases: AtomicUsize::new(0),
})
}
fn close_detached(&self) {
if !crate::rt::can_spawn() {
return;
}
let (endpoint, gossip) = (self.endpoint.clone(), self.gossip.clone());
crate::rt::spawn(async move {
let _ = crate::rt::time::timeout(std::time::Duration::from_secs(2), gossip.shutdown()).await;
let _ = crate::rt::time::timeout(std::time::Duration::from_secs(4), endpoint.close()).await;
});
}
pub fn endpoint(&self) -> &Endpoint {
&self.endpoint
}
pub fn node_key(&self) -> [u8; NODE_KEY_LEN] {
self.key
}
pub fn node_addr(&self) -> EndpointAddr {
wire::relay_only(self.endpoint.addr())
}
pub async fn neighbor_count(&self, topic: &[u8; 32]) -> usize {
self.topics
.read()
.await
.get(&TopicId::from_bytes(*topic))
.map(|t| t.neighbors.load(Ordering::Relaxed))
.unwrap_or(0)
}
pub async fn subscribe(&self, topic: [u8; 32], peers: Vec<EndpointAddr>, sink: mpsc::Sender<MeshEvent>) -> Result<(), String> {
let id = TopicId::from_bytes(topic);
let mut topics = loop {
let topics = self.topics.write().await;
if topics.contains_key(&id) {
return Err("already subscribed to this topic".into());
}
let pending = self.closing.lock().unwrap_or_else(|e| e.into_inner()).get(&id).cloned();
match pending {
None => break topics,
Some(gone) => {
drop(topics);
self.await_teardown(id, gone).await;
}
}
};
let peers: Vec<EndpointAddr> = peers.into_iter().map(|p| self.reachable(p)).collect();
let bootstrap: Vec<EndpointId> = peers.iter().map(|p| p.id).collect();
let sub = self
.gossip
.subscribe_with_opts(id, JoinOptions::with_bootstrap(bootstrap))
.await
.map_err(|e| e.to_string())?;
for peer in peers {
self.dial(peer);
}
let (sender, receiver) = sub.split();
let neighbors = Arc::new(AtomicUsize::new(0));
let pump = crate::rt::spawn(pump(receiver, sender.clone(), sink, self.key, neighbors.clone()));
let seq = AtomicI32::new(rand::random::<i32>());
topics.insert(id, Topic { sender, pump, seq, neighbors });
Ok(())
}
fn reachable(&self, peer: EndpointAddr) -> EndpointAddr {
let mut addr = wire::dialable(peer);
if addr.addrs.is_empty() {
if let Some(relay) = self.endpoint.addr().addrs.into_iter().find(|a| matches!(a, TransportAddr::Relay(_))) {
addr.addrs.insert(relay);
}
}
addr
}
fn dial(&self, peer: EndpointAddr) {
let peer = self.reachable(peer);
if peer.addrs.is_empty() {
return;
}
let (ep, g) = (self.endpoint.clone(), self.gossip.clone());
crate::rt::spawn(async move {
match ep.connect(peer, GOSSIP_ALPN).await {
Ok(conn) => {
if let Err(e) = g.handle_connection(conn).await {
crate::log_debug!("[xdc] peer connection ended: {e}");
}
}
Err(e) => crate::log_debug!("[xdc] dial failed: {e}"),
}
});
}
pub async fn add_peer(&self, topic: &[u8; 32], peer: EndpointAddr) -> Result<(), String> {
let mut last = String::new();
for attempt in 0..3u32 {
if attempt > 0 {
crate::rt::time::sleep(std::time::Duration::from_secs(1 << (attempt - 1))).await;
}
match crate::rt::time::timeout(ADD_PEER_TIMEOUT, self.try_add_peer(topic, &peer)).await {
Ok(Ok(())) => return Ok(()),
Ok(Err(e)) => last = e,
Err(_) => last = "timed out".into(),
}
}
Err(last)
}
async fn try_add_peer(&self, topic: &[u8; 32], peer: &EndpointAddr) -> Result<(), String> {
let addr = self.reachable(peer.clone());
let conn = self.endpoint.connect(addr, GOSSIP_ALPN).await.map_err(|e| e.to_string())?;
self.gossip.handle_connection(conn).await.map_err(|e| e.to_string())?;
let sender = {
let topics = self.topics.read().await;
topics.get(&TopicId::from_bytes(*topic)).ok_or("not subscribed")?.sender.clone()
};
sender.join_peers(vec![peer.id]).await.map_err(|e| e.to_string())
}
pub async fn broadcast(&self, topic: &[u8; 32], payload: Vec<u8>) -> Result<(), String> {
if payload.len() > wire::MAX_PAYLOAD {
return Err(format!("frame of {} bytes exceeds the {} byte limit", payload.len(), wire::MAX_PAYLOAD));
}
let (sender, frame) = {
let topics = self.topics.read().await;
let t = topics.get(&TopicId::from_bytes(*topic)).ok_or("not subscribed")?;
let seq = t.seq.fetch_add(1, Ordering::Relaxed).wrapping_add(1);
(t.sender.clone(), wire::seal_frame(payload, seq, &self.key))
};
sender.broadcast(frame.into()).await.map_err(|e| e.to_string())
}
pub async fn unsubscribe(&self, topic: &[u8; 32]) {
let id = TopicId::from_bytes(*topic);
let gone = {
let mut topics = self.topics.write().await;
match topics.remove(&id) {
Some(t) => {
let (done, gone) = watch::channel(false);
self.closing.lock().unwrap_or_else(|e| e.into_inner()).insert(id, gone.clone());
let closing = self.closing.clone();
crate::rt::spawn(async move {
drop(t.sender);
t.pump.abort();
let _ = t.pump.await;
crate::rt::yield_now().await;
closing.lock().unwrap_or_else(|e| e.into_inner()).remove(&id);
done.send_replace(true);
});
Some(gone)
}
None => self.closing.lock().unwrap_or_else(|e| e.into_inner()).get(&id).cloned(),
}
};
if let Some(gone) = gone {
self.await_teardown(id, gone).await;
}
}
async fn await_teardown(&self, id: TopicId, mut gone: watch::Receiver<bool>) {
if gone.wait_for(|g| *g).await.is_err() {
let mut closing = self.closing.lock().unwrap_or_else(|e| e.into_inner());
if closing.get(&id).is_some_and(|rx| rx.same_channel(&gone)) {
closing.remove(&id);
}
}
}
}
impl Drop for Mesh {
fn drop(&mut self) {
self.close_detached();
}
}
async fn pump(
mut receiver: GossipReceiver,
sender: GossipSender,
sink: mpsc::Sender<MeshEvent>,
our_key: [u8; NODE_KEY_LEN],
neighbors: Arc<AtomicUsize>,
) {
let lagged = AtomicBool::new(false);
let deliver = |ev: MeshEvent| {
if lagged.swap(false, Ordering::Relaxed) && sink.try_send(MeshEvent::Lagged).is_err() {
lagged.store(true, Ordering::Relaxed);
}
if sink.try_send(ev).is_err() {
lagged.store(true, Ordering::Relaxed);
}
};
while let Some(event) = receiver.next().await {
if sink.is_closed() {
break;
}
match event {
Ok(Event::Received(msg)) => {
let Some((payload, sender_key)) = wire::open_frame(&msg.content) else { continue };
if sender_key == our_key {
continue;
}
let direct = msg.scope.is_direct() && *msg.delivered_from.as_bytes() == sender_key;
deliver(MeshEvent::Data { payload: payload.to_vec(), sender: sender_key, direct });
}
Ok(Event::NeighborUp(peer)) => {
let _ = sender.join_peers(vec![peer]).await;
neighbors.fetch_add(1, Ordering::Relaxed);
deliver(MeshEvent::NeighborUp(*peer.as_bytes()));
}
Ok(Event::NeighborDown(peer)) => {
let _ = neighbors.fetch_update(Ordering::Relaxed, Ordering::Relaxed, |c| c.checked_sub(1));
deliver(MeshEvent::NeighborDown(*peer.as_bytes()));
}
Ok(Event::Lagged) => deliver(MeshEvent::Lagged),
Err(e) => crate::log_debug!("[xdc] gossip error: {e}"),
}
}
}
pub(crate) fn short(key: &[u8; NODE_KEY_LEN]) -> String {
crate::simd::hex::bytes_to_hex_string(&key[..6])
}
mod anyhow_lite {
pub type Result<T> = std::result::Result<T, Box<dyn std::error::Error + Send + Sync>>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn an_account_never_gets_another_accounts_mesh() {
let mut slot = Some((1u64, "a"));
assert!(matches!(claim(&mut slot, 1, true), Claim::Reuse("a")));
assert!(matches!(claim(&mut slot, 1, false), Claim::Reuse("a")), "its own, even mid-swap");
assert!(matches!(claim(&mut slot, 2, false), Claim::Refuse), "a swapped-out account binds nothing");
assert_eq!(slot, Some((1, "a")));
assert!(matches!(claim(&mut slot, 2, true), Claim::Bind(Some("a"))), "the next account replaces it");
assert_eq!(slot, None);
assert!(matches!(claim(&mut slot, 2, true), Claim::Bind(None)));
}
#[test]
fn retiring_closes_only_the_named_accounts_mesh() {
let mut slot = Some((2u64, "b"));
assert_eq!(take_owned(&mut slot, 1), None);
assert_eq!(slot, Some((2, "b")));
assert_eq!(take_owned(&mut slot, 2), Some("b"));
assert_eq!(slot, None);
}
#[tokio::test]
async fn an_unsubscribe_cut_off_half_way_still_drops_the_topic_before_a_retry_returns() {
let mesh = Mesh::offline().await;
let topic = [7u8; 32];
let (tx, mut rx) = mpsc::channel(4);
mesh.subscribe(topic, Vec::new(), tx).await.unwrap();
assert!(futures_util::FutureExt::now_or_never(mesh.unsubscribe(&topic)).is_none());
mesh.unsubscribe(&topic).await;
assert!(
matches!(rx.try_recv(), Err(mpsc::error::TryRecvError::Disconnected)),
"the pump's handles are gone once the retry returns"
);
let (tx, _rx) = mpsc::channel(4);
mesh.subscribe(topic, Vec::new(), tx).await.expect("the topic can be joined again");
mesh.unsubscribe(&topic).await;
}
#[tokio::test]
async fn releasing_a_mesh_closes_its_node() {
let mesh = Mesh::offline().await;
let endpoint = mesh.endpoint.clone();
drop(mesh);
for _ in 0..100 {
if endpoint.is_closed() {
return;
}
tokio::time::sleep(std::time::Duration::from_millis(20)).await;
}
panic!("the endpoint outlived its mesh");
}
#[tokio::test]
async fn a_mesh_is_retired_only_once_nothing_uses_it() {
let _serial = SLOT_TESTS.lock().await;
crate::db::with_session(crate::db::current_session(), async {
let owner = crate::db::current_session().id();
let mesh = install_for_test(owner, Mesh::offline().await).await;
let lease = lease().await.expect("the caller's own mesh");
assert!(Arc::ptr_eq(lease.mesh(), &mesh));
assert!(!retire_if_idle(&mesh).await, "a join between handing out the mesh and subscribing holds it");
drop(lease);
let (tx, _rx) = mpsc::channel(4);
mesh.subscribe([9u8; 32], Vec::new(), tx).await.unwrap();
assert!(!retire_if_idle(&mesh).await, "a subscribed topic holds it");
mesh.unsubscribe(&[9u8; 32]).await;
assert!(retire_if_idle(&mesh).await, "idle, it closes");
assert_eq!(slot_owner().await, None);
let successor = install_for_test(owner, Mesh::offline().await).await;
assert!(!retire_if_idle(&mesh).await, "a stale close must not take its successor");
assert_eq!(slot_owner().await, Some(owner));
assert!(retire_if_idle(&successor).await);
})
.await;
}
}