diff --git a/.containerignore b/.containerignore new file mode 100644 index 00000000000..3ea08526b04 --- /dev/null +++ b/.containerignore @@ -0,0 +1,2 @@ +target/ +.git/ diff --git a/iroh/src/socket/mapped_addrs.rs b/iroh/src/socket/mapped_addrs.rs index 33d682f748d..595cb311656 100644 --- a/iroh/src/socket/mapped_addrs.rs +++ b/iroh/src/socket/mapped_addrs.rs @@ -364,6 +364,28 @@ where let inner = self.inner.lock().expect("poisoned"); inner.lookup.get(addr).cloned() } + + pub(super) fn remove(&self, key: &K) -> Option { + let mut inner = self.inner.lock().expect("poisoned"); + if let Some(v) = inner.addrs.remove(key) { + inner.lookup.remove(&v); + Some(v) + } else { + None + } + } + + pub(super) fn retain(&self, mut f: impl FnMut(&K, &V) -> bool) { + let mut inner = self.inner.lock().expect("poisoned"); + let AddrMapInner { addrs, lookup } = &mut *inner; + addrs.retain(|k, v| { + let keep = f(k, v); + if !keep { + lookup.remove(v); + } + keep + }); + } } #[derive(Debug)] @@ -381,3 +403,46 @@ impl Default for AddrMapInner { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn remove_clears_forward_and_reverse_maps() { + let map: AddrMap = AddrMap::default(); + let addr = map.get(&1); + assert_eq!( + map.lookup(&addr), + Some(1), + "reverse lookup present after insert" + ); + + assert_eq!(map.remove(&1), Some(addr), "remove returns the mapped addr"); + assert_eq!(map.lookup(&addr), None, "reverse lookup cleared on remove"); + + // Re-inserting generates a fresh addr, proving the old entry was gone. + let addr2 = map.get(&1); + assert_ne!(addr, addr2); + } + + #[test] + fn retain_drops_unmatched_entries_and_their_reverse() { + let map: AddrMap = AddrMap::default(); + let kept = map.get(&1); + let dropped = map.get(&2); + + map.retain(|k, _| *k == 1); + + assert_eq!( + map.lookup(&kept), + Some(1), + "retained entry keeps its reverse lookup" + ); + assert_eq!( + map.lookup(&dropped), + None, + "dropped entry reverse lookup is cleared" + ); + } +} diff --git a/iroh/src/socket/remote_map.rs b/iroh/src/socket/remote_map.rs index 37c872cf0d7..db82f3874ab 100644 --- a/iroh/src/socket/remote_map.rs +++ b/iroh/src/socket/remote_map.rs @@ -236,7 +236,11 @@ impl RemoteMap { if leftover_msgs.is_empty() { // the actor shut down cleanly self.senders.remove(&remote_id); - trace!(%remote_id, "cleaned up RemoteStateActor"); + self.mapped_addrs.endpoint_addrs.remove(&remote_id); + self.mapped_addrs + .relay_addrs + .retain(|k, _| k.1 != remote_id); + trace!(%remote_id, "cleaned up RemoteStateActor and mapped_addrs"); true } else { // The remote actor got messages while it was closing, so we're restarting