From 875e6e124c50bdcbbe80bc569da77921d74fbd0e Mon Sep 17 00:00:00 2001 From: Mollusk Date: Fri, 29 May 2026 15:48:43 -0400 Subject: [PATCH] feat: redesign audio networking for real-network resilience Replaces the fire-and-forget datagram path with sequenced packets, a per-peer jitter buffer, and persistent per-peer send tasks. Together these fix three intertwined weaknesses that only showed up off localhost. Packet format: every audio frame now carries a 4-byte little-endian sequence number header ([seq][opus payload]), the basis for reordering and loss detection. Jitter buffer (core/jitter.rs): incoming packets are reordered by sequence behind a fixed ~60ms playout delay. Missing sequences with later packets already buffered are concealed via Opus PLC (decode(None)) -- a path the decoder supported but nothing ever invoked. Underruns go idle and re-buffer rather than concealing indefinitely. Covered by unit tests using real encoded frames (reorder, gap-conceal, prime, late-drop). Transport (network/iroh_impl.rs): each peer gets one long-lived send task fed by a shallow bounded channel (drop-oldest on backpressure), instead of spawning a throwaway task per peer per 20ms frame. Connections are now established reactively on peer-join and torn down on peer-leave; the lexicographically-lower EndpointId dials so a full-mesh pair forms exactly one shared bidirectional connection instead of two racing ones. This also removes the previous lock-held-across-connect().await serialization. Opus decoder: PLC output is now sized to one 20ms frame, so concealment synthesizes 20ms instead of a 120ms burst from the oversized max buffer. Known follow-up (Tier 2): no reconnect on transient connection loss; a send error currently retires the peer until they rejoin. Co-Authored-By: Claude Opus 4.8 --- src/codec/opus_impl.rs | 48 +++++---- src/core/jitter.rs | 192 ++++++++++++++++++++++++++++++++++++ src/core/mod.rs | 152 ++++++++++++++-------------- src/network/iroh_impl.rs | 208 +++++++++++++++++++++++++-------------- src/network/mod.rs | 13 ++- 5 files changed, 437 insertions(+), 176 deletions(-) create mode 100644 src/core/jitter.rs diff --git a/src/codec/opus_impl.rs b/src/codec/opus_impl.rs index 93afe52..b701f3e 100644 --- a/src/codec/opus_impl.rs +++ b/src/codec/opus_impl.rs @@ -31,44 +31,52 @@ impl AudioEncoder for OpusEncoder { pub struct OpusDecoder { decoder: Decoder, channels: Channels, + /// Samples-per-channel of the frames we transmit (20ms @ 48kHz = 960). + /// Used to size the Packet Loss Concealment output, since libopus conceals + /// `frame_size` samples when given no input — passing the full max buffer + /// would synthesize a 120ms burst instead of a single 20ms frame. + frame_samples: usize, } impl OpusDecoder { /// Creates a new Opus decoder. - /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono - pub fn new(sample_rate: u32, channels: Channels) -> Result { + /// Standard voice parameters: sample_rate = 48000, channels = Channels::Mono. + /// `frame_samples` is the per-channel length of one transmitted frame (e.g. 960). + pub fn new(sample_rate: u32, channels: Channels, frame_samples: usize) -> Result { let decoder = Decoder::new(sample_rate, channels) .map_err(|e| CodecError::Init(format!("Failed to create Opus decoder: {}", e)))?; - Ok(Self { decoder, channels }) + Ok(Self { decoder, channels, frame_samples }) + } + + fn channels_count(&self) -> usize { + match self.channels { + Channels::Mono => 1, + Channels::Stereo => 2, + } } } impl AudioDecoder for OpusDecoder { fn decode(&mut self, compressed: Option<&[u8]>) -> Result, CodecError> { - // Maximum Opus frame size is 120ms. At 48kHz, this is 5760 samples per channel. - let channels_count = match self.channels { - Channels::Mono => 1, - Channels::Stereo => 2, - }; - let max_samples = 5760 * channels_count; - let mut pcm = vec![0i16; max_samples]; + let channels_count = self.channels_count(); - let decoded_samples_per_channel = match compressed { + let (mut pcm, input): (Vec, &[u8]) = match compressed { Some(data) if !data.is_empty() => { - // Normal decode - self.decoder.decode(data, &mut pcm, false) - .map_err(|e| CodecError::Decode(format!("Opus decoding failed: {}", e)))? + // Normal decode. Size the buffer to the maximum Opus frame (120ms = + // 5760 samples/channel); libopus decodes the packet's true duration. + (vec![0i16; 5760 * channels_count], data) } _ => { - // Packet Loss Concealment (PLC) - // In opus-rs, passing an empty slice triggers PLC. - self.decoder.decode(&[], &mut pcm, false) - .map_err(|e| CodecError::Decode(format!("Opus PLC decoding failed: {}", e)))? + // Packet Loss Concealment: an empty input makes libopus synthesize + // exactly `frame_samples` of concealment, so size the buffer to match. + (vec![0i16; self.frame_samples * channels_count], &[]) } }; - let total_samples = decoded_samples_per_channel * channels_count; - pcm.truncate(total_samples); + let decoded_per_channel = self.decoder.decode(input, &mut pcm, false) + .map_err(|e| CodecError::Decode(format!("Opus decoding failed: {}", e)))?; + + pcm.truncate(decoded_per_channel * channels_count); Ok(pcm) } } diff --git a/src/core/jitter.rs b/src/core/jitter.rs new file mode 100644 index 0000000..72ac0f1 --- /dev/null +++ b/src/core/jitter.rs @@ -0,0 +1,192 @@ +//! Per-peer jitter buffer with Opus packet-loss concealment. +//! +//! Incoming audio arrives as unreliable QUIC datagrams that can be reordered, +//! duplicated, or dropped on real networks. Each packet carries a monotonic +//! sequence number (assigned by the sender). This buffer reorders packets by +//! sequence, holds a small fixed playout delay to absorb jitter, and — when a +//! sequence is missing but later packets have already arrived — synthesizes a +//! concealment frame via Opus PLC instead of emitting a click of silence. + +use crate::codec::{AudioDecoder, CodecError, opus_impl::OpusDecoder}; +use opus::Channels; +use std::collections::BTreeMap; + +/// Samples per channel in one transmitted frame (20ms @ 48kHz mono). +pub const FRAME_SAMPLES: usize = 960; + +/// How many frames to buffer before playout begins (~60ms). This is the +/// tolerance window for reordering and jitter; larger = more resilient but +/// more latency. +const TARGET_DELAY_FRAMES: usize = 3; + +/// Hard cap on buffered frames (~640ms). If we ever exceed this we've fallen +/// badly behind, so we drop the oldest and resync rather than grow unbounded. +const MAX_BUFFERED_FRAMES: usize = 32; + +pub struct JitterBuffer { + decoder: OpusDecoder, + /// Reorder window: sequence number -> encoded Opus payload. + packets: BTreeMap>, + /// Next sequence we expect to play. `None` means idle/buffering: we are + /// waiting to accumulate `TARGET_DELAY_FRAMES` before (re)starting playout. + next_seq: Option, +} + +/// Wrapping-aware "is `a` strictly before `b`" for sequence numbers. +fn seq_before(a: u32, b: u32) -> bool { + a != b && b.wrapping_sub(a) < (1 << 31) +} + +impl JitterBuffer { + pub fn new() -> Result { + Ok(Self { + decoder: OpusDecoder::new(48000, Channels::Mono, FRAME_SAMPLES)?, + packets: BTreeMap::new(), + next_seq: None, + }) + } + + /// Store a received packet, dropping ones we've already played past and + /// bounding total depth. + pub fn insert(&mut self, seq: u32, payload: Vec) { + // Too late: this sequence has already been played (or concealed). + if let Some(next) = self.next_seq + && seq_before(seq, next) + { + return; + } + self.packets.insert(seq, payload); + + while self.packets.len() > MAX_BUFFERED_FRAMES { + let oldest = *self.packets.keys().next().expect("non-empty"); + self.packets.remove(&oldest); + // We've discarded backlog; resync the playout head to the new front. + self.next_seq = self.packets.keys().next().copied(); + } + } + + /// Produce the next 20ms PCM frame for playout, or `None` when idle or + /// still buffering (the caller should treat `None` as silence). + pub fn pop_frame(&mut self) -> Option> { + match self.next_seq { + None => { + // Buffering: start playout once we have enough to absorb jitter. + if self.packets.len() >= TARGET_DELAY_FRAMES { + self.next_seq = self.packets.keys().next().copied(); + self.pop_frame() + } else { + None + } + } + Some(next) => { + if let Some(payload) = self.packets.remove(&next) { + self.next_seq = Some(next.wrapping_add(1)); + self.decoder.decode(Some(&payload)).ok() + } else if self.packets.is_empty() { + // Underrun: the talker has gone quiet (or stopped). Go idle + // and re-buffer before resuming, rather than concealing forever. + self.next_seq = None; + None + } else { + // Gap with later packets already buffered: a packet was lost + // or reordered out of window. Conceal this frame via Opus PLC. + self.next_seq = Some(next.wrapping_add(1)); + self.decoder.decode(None).ok() + } + } + } + } + + /// True when nothing is buffered and playout is idle (talker silent). + pub fn is_idle(&self) -> bool { + self.next_seq.is_none() && self.packets.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::codec::AudioEncoder; + use crate::codec::opus_impl::OpusEncoder; + use opus::{Application, Channels}; + + /// A real, decodable Opus packet for one 20ms mono frame at amplitude `amp`. + fn frame(enc: &mut OpusEncoder, amp: i16) -> Vec { + let pcm: Vec = (0..FRAME_SAMPLES) + .map(|i| if i % 2 == 0 { amp } else { -amp }) + .collect(); + enc.encode(&pcm).unwrap() + } + + #[test] + fn buffers_then_plays_in_order() { + let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap(); + let mut jb = JitterBuffer::new().unwrap(); + + // Below the target delay, playout hasn't primed yet. + jb.insert(0, frame(&mut enc, 1000)); + assert!(jb.pop_frame().is_none()); + + // Reaching the target delay primes playout and yields the first frame. + jb.insert(1, frame(&mut enc, 1000)); + jb.insert(2, frame(&mut enc, 1000)); + assert_eq!(jb.pop_frame().map(|f| f.len()), Some(FRAME_SAMPLES)); + assert_eq!(jb.pop_frame().map(|f| f.len()), Some(FRAME_SAMPLES)); + assert_eq!(jb.pop_frame().map(|f| f.len()), Some(FRAME_SAMPLES)); + // Drained: idle again. + assert!(jb.pop_frame().is_none()); + assert!(jb.is_idle()); + } + + #[test] + fn reorders_out_of_order_arrivals() { + let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap(); + let mut jb = JitterBuffer::new().unwrap(); + + // Arrive scrambled but within the buffering window. + jb.insert(2, frame(&mut enc, 800)); + jb.insert(0, frame(&mut enc, 800)); + jb.insert(1, frame(&mut enc, 800)); + + // Three real frames come out (in sequence order), then idle. + assert!(jb.pop_frame().is_some()); + assert!(jb.pop_frame().is_some()); + assert!(jb.pop_frame().is_some()); + assert!(jb.pop_frame().is_none()); + } + + #[test] + fn conceals_gap_when_later_packets_present() { + let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap(); + let mut jb = JitterBuffer::new().unwrap(); + + // Seq 2 is missing, but 0,1,3 arrive — enough to prime. + jb.insert(0, frame(&mut enc, 1200)); + jb.insert(1, frame(&mut enc, 1200)); + jb.insert(3, frame(&mut enc, 1200)); + + assert!(jb.pop_frame().is_some()); // seq 0 + assert!(jb.pop_frame().is_some()); // seq 1 + // seq 2 missing but seq 3 buffered -> Opus PLC produces a concealment frame. + let concealed = jb.pop_frame(); + assert_eq!(concealed.map(|f| f.len()), Some(FRAME_SAMPLES)); + assert!(jb.pop_frame().is_some()); // seq 3 + assert!(jb.pop_frame().is_none()); + } + + #[test] + fn drops_packets_already_played() { + let mut enc = OpusEncoder::new(48000, Channels::Mono, Application::Voip).unwrap(); + let mut jb = JitterBuffer::new().unwrap(); + + jb.insert(5, frame(&mut enc, 600)); + jb.insert(6, frame(&mut enc, 600)); + jb.insert(7, frame(&mut enc, 600)); + assert!(jb.pop_frame().is_some()); // primes at seq 5, plays 5 + assert!(jb.pop_frame().is_some()); // 6 + + // A straggler for an already-played sequence must be discarded. + jb.insert(5, frame(&mut enc, 600)); + assert_eq!(jb.packets.len(), 1); // only seq 7 remains buffered + } +} diff --git a/src/core/mod.rs b/src/core/mod.rs index b1f6c45..4a3df83 100644 --- a/src/core/mod.rs +++ b/src/core/mod.rs @@ -1,7 +1,9 @@ pub mod messages; +pub mod jitter; use crate::audio::{AudioBackend, pipewire_impl::PipeWireBackend}; -use crate::codec::{AudioEncoder, AudioDecoder, opus_impl::{OpusEncoder, OpusDecoder}}; +use crate::codec::{AudioEncoder, opus_impl::OpusEncoder}; +use crate::core::jitter::{JitterBuffer, FRAME_SAMPLES}; use crate::network::{ NetworkTransport, RoomState, PeerState, RoomEvent, PeerSpeakTicket, iroh_impl::IrohTransport, @@ -12,7 +14,7 @@ use crate::core::messages::{CoreCommand, UiEvent}; use iroh::{Endpoint, EndpointId, endpoint::presets, protocol::Router}; use iroh_gossip::net::Gossip; use tokio::sync::{mpsc, Mutex}; -use std::collections::{HashMap, VecDeque}; +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::time::Duration; @@ -189,7 +191,7 @@ async fn run_core_loop( continue; } - let queues: Arc>>> = Arc::new(Mutex::new(HashMap::new())); + let jitter: Arc>> = Arc::new(Mutex::new(HashMap::new())); // 1. Capture & encoding thread let is_muted_clone = is_muted.clone(); @@ -197,8 +199,6 @@ async fn run_core_loop( let ptt_active_clone = ptt_active.clone(); let noise_gate_threshold_clone = noise_gate_threshold.clone(); let transport_clone = transport.clone(); - let room_state_clone = room_state.clone(); - let tokio_handle = tokio::runtime::Handle::current(); let capture_thread = std::thread::spawn(move || { use opus::{Channels, Application}; @@ -209,6 +209,9 @@ async fn run_core_loop( return; } }; + // Per-sender packet sequence number, prepended to every frame so + // receivers can reorder and conceal loss. Wraps after ~years. + let mut seq: u32 = 0; while let Ok(pcm) = capture_rx.recv() { if is_muted_clone.load(Ordering::Relaxed) { @@ -217,7 +220,7 @@ async fn run_core_loop( if ptt_mode_clone.load(Ordering::Relaxed) && !ptt_active_clone.load(Ordering::Relaxed) { continue; } - + let ng_bits = noise_gate_threshold_clone.load(Ordering::Relaxed); let ng_thresh = f32::from_bits(ng_bits); if ng_thresh > 0.0001 { @@ -233,26 +236,22 @@ async fn run_core_loop( } if let Ok(encoded) = encoder.encode(&pcm) { - let bytes = bytes::Bytes::from(encoded); - let active = room_state_clone.active_peers(); - for (peer_id, _) in active { - let transport = transport_clone.clone(); - let bytes = bytes.clone(); - tokio_handle.spawn(async move { - if let Err(e) = transport.send_datagram(peer_id, bytes).await { - crate::log_msg(&format!("Failed to send datagram to peer {:?}: {:?}", peer_id, e)); - } - }); - } + // Frame on the wire: [seq: u32 LE][opus payload]. + let mut packet = Vec::with_capacity(4 + encoded.len()); + packet.extend_from_slice(&seq.to_le_bytes()); + packet.extend_from_slice(&encoded); + seq = seq.wrapping_add(1); + transport_clone.broadcast(bytes::Bytes::from(packet)); } } }); - // 2. Receiver & decoding task + // 2. Receiver task: parse the sequence header and hand each packet + // to that peer's jitter buffer. Decoding happens later, on the + // playout side, so loss can be concealed at the right moment. let transport_recv = transport.clone(); - let queues_recv = queues.clone(); + let jitter_recv = jitter.clone(); let datagram_task = tokio::spawn(async move { - use opus::Channels; let mut datagram_rx = match transport_recv.receive_datagrams().await { Ok(rx) => rx, Err(e) => { @@ -261,37 +260,34 @@ async fn run_core_loop( } }; - let mut decoders: HashMap = HashMap::new(); while let Some((from_peer, bytes)) = datagram_rx.recv().await { - crate::log_msg(&format!("Received datagram from peer={:?}, len={}", from_peer, bytes.len())); - let decoder = match decoders.entry(from_peer) { + if bytes.len() < 4 { + continue; // malformed: missing sequence header + } + let seq = u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + let payload = bytes[4..].to_vec(); + + let mut guard = jitter_recv.lock().await; + let buffer = match guard.entry(from_peer) { std::collections::hash_map::Entry::Occupied(entry) => entry.into_mut(), std::collections::hash_map::Entry::Vacant(entry) => { - match OpusDecoder::new(48000, Channels::Mono) { - Ok(dec) => entry.insert(dec), + match JitterBuffer::new() { + Ok(jb) => entry.insert(jb), Err(e) => { - crate::log_msg(&format!("Failed to initialize decoder for {:?}: {:?}", from_peer, e)); + crate::log_msg(&format!("Failed to init jitter buffer for {:?}: {:?}", from_peer, e)); continue; } } } }; - - match decoder.decode(Some(&bytes)) { - Ok(pcm) => { - let mut guard = queues_recv.lock().await; - let queue = guard.entry(from_peer).or_insert_with(VecDeque::new); - queue.extend(pcm); - } - Err(e) => { - crate::log_msg(&format!("Failed to decode packet from {:?}: {:?}", from_peer, e)); - } - } + buffer.insert(seq, payload); } }); - // 3. Mixing & level extraction loop task - let queues_mixer = queues.clone(); + // 3. Mixing & level extraction loop task. Every 20ms, pull one + // concealed frame per peer from its jitter buffer, apply + // per-peer volume, sum, and hand the mix to playback. + let jitter_mixer = jitter.clone(); let is_deafened_clone = is_deafened.clone(); let peer_volumes_mixer = peer_volumes.clone(); let ui_tx_mixer = ui_tx.clone(); @@ -302,55 +298,49 @@ async fn run_core_loop( loop { interval.tick().await; - let mut guard = queues_mixer.lock().await; - let mut mixed = vec![0i16; 960]; + let current_volumes = peer_volumes_mixer.lock().await.clone(); let mut active_levels = Vec::new(); let mut peer_frames = Vec::new(); - let current_volumes = peer_volumes_mixer.lock().await.clone(); - for (&peer_id, queue) in guard.iter_mut() { - let mut frame = vec![0i16; 960]; - let len = queue.len(); - if len >= 960 { - if len > 9600 { - let drain = len - 960; - queue.drain(0..drain); - } - for sample in frame.iter_mut() { - *sample = queue.pop_front().unwrap_or(0); - } - } else { - for sample in frame.iter_mut().take(len) { - *sample = queue.pop_front().unwrap_or(0); + { + let mut guard = jitter_mixer.lock().await; + for (&peer_id, buffer) in guard.iter_mut() { + // `None` means idle/buffering: contribute nothing + // and report a zero level so the UI shows idle. + let Some(mut frame) = buffer.pop_frame() else { + active_levels.push((peer_id, 0.0)); + continue; + }; + + let vol = current_volumes.get(&peer_id).copied().unwrap_or(1.0); + if (vol - 1.0).abs() > f32::EPSILON { + for sample in frame.iter_mut() { + *sample = (*sample as f32 * vol).clamp(i16::MIN as f32, i16::MAX as f32) as i16; + } } + + let sum_sq: f32 = frame.iter().map(|&x| (x as f32).powi(2)).sum(); + let rms = (sum_sq / frame.len().max(1) as f32).sqrt(); + let level = (rms / 32768.0).clamp(0.0, 1.0); + active_levels.push((peer_id, level)); + + peer_frames.push(frame); } - - let vol = current_volumes.get(&peer_id).copied().unwrap_or(1.0); - for sample in frame.iter_mut() { - *sample = (*sample as f32 * vol).clamp(i16::MIN as f32, i16::MAX as f32) as i16; - } - - // Calculate speaking level (RMS normalized) - let sum_sq: f32 = frame.iter().map(|&x| (x as f32).powi(2)).sum(); - let rms = (sum_sq / 960.0).sqrt(); - let level = (rms / 32768.0).clamp(0.0, 1.0); - active_levels.push((peer_id, level)); - - peer_frames.push(frame); } + let mut mixed = vec![0i16; FRAME_SAMPLES]; if !peer_frames.is_empty() { - for i in 0..960 { - let mut sum = 0i32; - for f in &peer_frames { - sum += f[i] as i32; - } - mixed[i] = sum.clamp(i16::MIN as i32, i16::MAX as i32) as i16; + for (i, out) in mixed.iter_mut().enumerate() { + let sum: i32 = peer_frames + .iter() + .map(|f| f.get(i).copied().unwrap_or(0) as i32) + .sum(); + *out = sum.clamp(i16::MIN as i32, i16::MAX as i32) as i16; } } let frame_to_send = if is_deafened_clone.load(Ordering::Relaxed) { - vec![0i16; 960] + vec![0i16; FRAME_SAMPLES] } else { mixed }; @@ -372,16 +362,20 @@ async fn run_core_loop( } }; let ui_tx_events = ui_tx.clone(); - let queues_events = queues.clone(); + let jitter_events = jitter.clone(); + let transport_events = transport.clone(); let event_task = tokio::spawn(async move { while let Some(event) = room_events.recv().await { match event { RoomEvent::PeerJoined(peer_id, state) => { - queues_events.lock().await.entry(peer_id).or_insert_with(VecDeque::new); + // Establish the audio connection as soon as the peer + // is known (the transport dedupes the full-mesh race). + transport_events.connect_peer(peer_id).await; let _ = ui_tx_events.send(UiEvent::PeerJoined { id: peer_id, state }).await; } RoomEvent::PeerLeft(peer_id) => { - queues_events.lock().await.remove(&peer_id); + transport_events.disconnect_peer(peer_id).await; + jitter_events.lock().await.remove(&peer_id); let _ = ui_tx_events.send(UiEvent::PeerLeft { id: peer_id }).await; } RoomEvent::PeerUpdated(peer_id, state) => { diff --git a/src/network/iroh_impl.rs b/src/network/iroh_impl.rs index 390004d..93d05c5 100644 --- a/src/network/iroh_impl.rs +++ b/src/network/iroh_impl.rs @@ -2,16 +2,98 @@ use crate::network::{NetworkTransport, NetError}; use iroh::{Endpoint, EndpointId}; use iroh::endpoint::Connection; use bytes::Bytes; -use tokio::sync::{mpsc, Mutex}; +use tokio::sync::mpsc; use tokio::sync::mpsc::Receiver; -use std::sync::Arc; +use std::sync::{Arc, Mutex as StdMutex}; use std::collections::HashMap; use async_trait::async_trait; -#[derive(Debug, Clone)] -pub struct AudioProtocol { +const AUDIO_ALPN: &[u8] = b"peerspeak-audio"; + +/// Per-peer datagram send queue depth. Audio is real-time, so a backlog is +/// useless latency — keep it shallow and drop the oldest frame when full. +const SEND_QUEUE_DEPTH: usize = 8; + +/// State shared between the transport and its protocol handler so both inbound +/// (accepted) and outbound (dialed) connections register the same way. +struct Shared { + /// Sync-lockable send handles, so `broadcast` can fan out from the (non-async) + /// capture/encode thread without touching the Tokio runtime. + senders: StdMutex>>, + /// Per-peer task handles + a retained connection clone. Keeping the clone + /// alive is what stops iroh from closing an accepted connection once the + /// `accept()` future returns. + peers: tokio::sync::Mutex>, incoming_tx: mpsc::Sender<(EndpointId, Bytes)>, - connections: Arc>>, +} + +struct PeerTasks { + send_task: tokio::task::JoinHandle<()>, + read_task: tokio::task::JoinHandle<()>, + _conn: Connection, +} + +impl Shared { + /// Register a live connection: spin up its send loop (datagrams out) and + /// read loop (datagrams in). Idempotent — a second registration for an + /// already-known peer is ignored so we never run duplicate loops. + async fn register(self: &Arc, peer_id: EndpointId, conn: Connection) { + { + let peers = self.peers.lock().await; + if peers.contains_key(&peer_id) { + return; + } + } + + let (send_tx, mut send_rx) = mpsc::channel::(SEND_QUEUE_DEPTH); + + let conn_send = conn.clone(); + let send_task = tokio::spawn(async move { + while let Some(data) = send_rx.recv().await { + if conn_send.send_datagram(data).is_err() { + break; + } + } + }); + + let conn_read = conn.clone(); + let incoming_tx = self.incoming_tx.clone(); + let read_task = tokio::spawn(async move { + while let Ok(bytes) = conn_read.read_datagram().await { + if incoming_tx.send((peer_id, bytes)).await.is_err() { + break; + } + } + }); + + self.senders.lock().unwrap().insert(peer_id, send_tx); + self.peers.lock().await.insert( + peer_id, + PeerTasks { send_task, read_task, _conn: conn }, + ); + crate::log_msg(&format!("Transport: registered peer {:?}", peer_id)); + } + + async fn remove(&self, peer_id: EndpointId) { + self.senders.lock().unwrap().remove(&peer_id); + if let Some(tasks) = self.peers.lock().await.remove(&peer_id) { + tasks.send_task.abort(); + tasks.read_task.abort(); + // Dropping `_conn` (the last retained clone) closes the connection. + crate::log_msg(&format!("Transport: removed peer {:?}", peer_id)); + } + } +} + +#[derive(Clone)] +pub struct AudioProtocol { + shared: Arc, +} + +impl std::fmt::Debug for AudioProtocol { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("AudioProtocol").finish_non_exhaustive() + } } impl iroh::protocol::ProtocolHandler for AudioProtocol { @@ -20,24 +102,11 @@ impl iroh::protocol::ProtocolHandler for AudioProtocol { connection: Connection, ) -> impl std::future::Future> + Send { let peer_id = connection.remote_id(); - let incoming_tx = self.incoming_tx.clone(); - let connections = self.connections.clone(); - + let shared = self.shared.clone(); async move { - connections.lock().await.insert(peer_id, connection.clone()); - loop { - match connection.read_datagram().await { - Ok(bytes) => { - if incoming_tx.send((peer_id, bytes)).await.is_err() { - break; - } - } - Err(_) => { - connections.lock().await.remove(&peer_id); - break; - } - } - } + // Register and return: the retained connection clone in `PeerTasks` + // keeps the connection open after this future resolves. + shared.register(peer_id, connection).await; Ok(()) } } @@ -45,81 +114,70 @@ impl iroh::protocol::ProtocolHandler for AudioProtocol { pub struct IrohTransport { endpoint: Endpoint, - connections: Arc>>, - incoming_tx: mpsc::Sender<(EndpointId, Bytes)>, - incoming_rx: Mutex>>, + self_id: EndpointId, + shared: Arc, + incoming_rx: tokio::sync::Mutex>>, } impl IrohTransport { pub fn new(endpoint: Endpoint) -> (Self, AudioProtocol) { let (incoming_tx, incoming_rx) = mpsc::channel(1000); - let connections = Arc::new(Mutex::new(HashMap::new())); + let self_id = endpoint.id(); - let audio_proto = AudioProtocol { - incoming_tx: incoming_tx.clone(), - connections: connections.clone(), - }; + let shared = Arc::new(Shared { + senders: StdMutex::new(HashMap::new()), + peers: tokio::sync::Mutex::new(HashMap::new()), + incoming_tx, + }); + let protocol = AudioProtocol { shared: shared.clone() }; let transport = Self { endpoint, - connections, - incoming_tx, - incoming_rx: Mutex::new(Some(incoming_rx)), + self_id, + shared, + incoming_rx: tokio::sync::Mutex::new(Some(incoming_rx)), }; - (transport, audio_proto) + (transport, protocol) } } #[async_trait] impl NetworkTransport for IrohTransport { - async fn send_datagram(&self, peer_id: EndpointId, data: Bytes) -> Result<(), NetError> { - let mut conns = self.connections.lock().await; - let conn = if let Some(conn) = conns.get(&peer_id) { - conn.clone() - } else { - // Establish a new connection. - // We use the same audio ALPN: b"peerspeak-audio" - let alpn = b"peerspeak-audio"; - let conn = self.endpoint.connect(peer_id, alpn).await - .map_err(|e| NetError::Connection(e.to_string()))?; + async fn connect_peer(&self, peer_id: EndpointId) { + // Deterministic initiator: only the lexicographically-lower id dials, so + // a full-mesh pair forms exactly one shared connection instead of two + // racing ones. The higher id waits for the inbound `accept()`. + if self.self_id.to_string() >= peer_id.to_string() { + return; + } + if self.shared.peers.lock().await.contains_key(&peer_id) { + return; + } - conns.insert(peer_id, conn.clone()); + match self.endpoint.connect(peer_id, AUDIO_ALPN).await { + Ok(conn) => self.shared.register(peer_id, conn).await, + Err(e) => crate::log_msg(&format!("Transport: dial to {:?} failed: {:?}", peer_id, e)), + } + } - let incoming_tx_inner = self.incoming_tx.clone(); - let connections_inner = self.connections.clone(); - let conn_clone = conn.clone(); + async fn disconnect_peer(&self, peer_id: EndpointId) { + self.shared.remove(peer_id).await; + } - tokio::spawn(async move { - loop { - match conn_clone.read_datagram().await { - Ok(bytes) => { - if incoming_tx_inner.send((peer_id, bytes)).await.is_err() { - break; - } - } - Err(_) => { - connections_inner.lock().await.remove(&peer_id); - break; - } - } - } - }); - - conn - }; - - conn.send_datagram(data) - .map_err(|e| NetError::Connection(e.to_string()))?; - Ok(()) + fn broadcast(&self, data: Bytes) { + let senders = self.shared.senders.lock().unwrap(); + for tx in senders.values() { + // Drop on a full queue: stale audio is worthless, and we must never + // block the encode thread on a slow peer. + let _ = tx.try_send(data.clone()); + } } async fn receive_datagrams(&self) -> Result, NetError> { let mut rx_guard = self.incoming_rx.lock().await; - if let Some(rx) = rx_guard.take() { - Ok(rx) - } else { - Err(NetError::Other("Datagram receiver already subscribed".to_string())) - } + rx_guard + .take() + .ok_or_else(|| NetError::Other("Datagram receiver already subscribed".to_string())) } } diff --git a/src/network/mod.rs b/src/network/mod.rs index 853f5d7..00eea11 100644 --- a/src/network/mod.rs +++ b/src/network/mod.rs @@ -69,8 +69,17 @@ impl FromStr for PeerSpeakTicket { #[async_trait] pub trait NetworkTransport: Send + Sync { - /// Send a low-latency unreliable datagram to a specific peer (for audio). - async fn send_datagram(&self, peer_id: EndpointId, data: Bytes) -> Result<(), NetError>; + /// Establish (or ensure) a connection to a peer and set up its send path. + /// Idempotent; safe to call again for an already-connected peer. + async fn connect_peer(&self, peer_id: EndpointId); + + /// Tear down the connection and send path for a peer that has left. + async fn disconnect_peer(&self, peer_id: EndpointId); + + /// Fan a single audio datagram out to every connected peer. Non-blocking: + /// per-peer queues drop the oldest-pending frame when full, so a slow link + /// can never stall the capture/encode thread. Callable from any thread. + fn broadcast(&self, data: Bytes); /// Subscribes to incoming datagrams from any peer. async fn receive_datagrams(&self) -> Result, NetError>;