refactor: use boxed watcher, not watchable, on connection (#3632)

## Description

Alternative to #3631

Replaces the `Watchable`s for path changes on the `Connection` with a
boxed `Watcher`. The watcher is boxed because it would increase the
`Connection` struct size significantly otherwise because the
mapped-and-joined watcher with a `SmallVec` of `PathInfo` inside is ~600
bytes atm.

The benefit of storing a `Watcher` and not a `Watchable` is that the
watcher streams now close once the EndpointStateActor drops the state
for the connection, which it does after the connection is closed.

Also adds a test for path watching, including testing that the streams
now close when the connection closes.

## Breaking Changes

<!-- Optional, if there are any breaking changes document them,
including how to migrate older code. -->

## Notes & open questions

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

## Change checklist
<!-- Remove any that are not relevant. -->
- [ ] 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.
- [ ] List all breaking changes in the above "Breaking Changes" section.
- [ ] Open an issue or PR on any number0 repos that are affected by this
breaking change. Give guidance on how the updates should be handled or
do the actual updates themselves. The major ones are:
    - [ ] [`quic-rpc`](https://github.com/n0-computer/quic-rpc)
    - [ ] [`iroh-gossip`](https://github.com/n0-computer/iroh-gossip)
    - [ ] [`iroh-blobs`](https://github.com/n0-computer/iroh-blobs)
    - [ ] [`dumbpipe`](https://github.com/n0-computer/dumbpipe)
    - [ ] [`sendme`](https://github.com/n0-computer/sendme)

Author: Frando

iroh/src/endpoint/connection.rs

@@ -43,7 +43,7 @@ use crate::{
     discovery::DiscoveryTask,
     magicsock::{
         EndpointStateActorStoppedError,
-        endpoint_map::{PathInfoList, PathsWatchable},
+        endpoint_map::{PathInfoList, PathsWatcher},
     },
 };
 
@@ -1213,7 +1213,7 @@ pub struct Connection {
     inner: quinn::Connection,
     remote_id: EndpointId,
     alpn: Vec<u8>,
-    paths: PathsWatchable,
+    paths: PathsWatcher,
 }
 
 #[allow(missing_docs)]
@@ -1463,7 +1463,7 @@ impl Connection {
     /// [`PathInfo::is_selected`]: crate::magicsock::PathInfo::is_selected
     /// [`PathInfo`]: crate::magicsock::PathInfo
     pub fn paths(&self) -> impl Watcher<Value = PathInfoList> + Unpin + Send + Sync + 'static {
-        self.paths.watch(self.inner.weak_handle())
+        self.paths.clone()
     }
 
     /// Derives keying material from this connection's TLS session secrets.
@@ -1511,16 +1511,21 @@ impl Connection {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use iroh_base::{EndpointAddr, SecretKey};
     use n0_error::{Result, StackResultExt, StdResultExt};
+    use n0_future::StreamExt;
+    use n0_watcher::Watcher;
     use rand::SeedableRng;
-    use tracing::{Instrument, info_span, trace_span};
+    use tracing::{Instrument, error_span, info_span, trace_span};
     use tracing_test::traced_test;
 
     use super::Endpoint;
     use crate::{
         RelayMode,
-        endpoint::{ConnectOptions, Incoming, ZeroRttStatus},
+        endpoint::{ConnectOptions, Incoming, PathInfo, PathInfoList, ZeroRttStatus},
+        test_utils::run_relay_server,
     };
 
     const TEST_ALPN: &[u8] = b"n0/iroh/test";
@@ -1730,4 +1735,94 @@ mod tests {
         tokio::join!(client.close(), server.close());
         Ok(())
     }
+
+    #[tokio::test]
+    async fn test_paths_watcher() -> Result {
+        tracing_subscriber::fmt::init();
+        const ALPN: &[u8] = b"test";
+        let mut rng = rand_chacha::ChaCha8Rng::seed_from_u64(0u64);
+        let (relay_map, _relay_map, _guard) = run_relay_server().await?;
+        let server = Endpoint::empty_builder(RelayMode::Custom(relay_map.clone()))
+            .secret_key(SecretKey::generate(&mut rng))
+            .insecure_skip_relay_cert_verify(true)
+            .alpns(vec![ALPN.to_vec()])
+            .bind()
+            .await?;
+
+        let client = Endpoint::empty_builder(RelayMode::Custom(relay_map.clone()))
+            .secret_key(SecretKey::generate(&mut rng))
+            .insecure_skip_relay_cert_verify(true)
+            .bind()
+            .await?;
+
+        server.online().await;
+        let server_addr = server.addr();
+        tracing::info!("server addr: {server_addr:?}");
+
+        let (conn_client, conn_server) = tokio::join!(
+            async { client.connect(server_addr, ALPN).await.unwrap() },
+            async { server.accept().await.unwrap().await.unwrap() }
+        );
+        tracing::info!("connected");
+        let mut paths_client = conn_client.paths().stream();
+        let mut paths_server = conn_server.paths().stream();
+
+        /// Advances the path stream until at least one IP and one relay paths are available.
+        ///
+        /// Panics if the path stream finishes before that happens.
+        async fn wait_for_paths(
+            stream: &mut n0_watcher::Stream<impl n0_watcher::Watcher<Value = PathInfoList> + Unpin>,
+        ) {
+            loop {
+                let paths = stream.next().await.expect("paths stream ended");
+                tracing::info!(?paths, "paths");
+                if paths.len() >= 2
+                    && paths.iter().any(PathInfo::is_relay)
+                    && paths.iter().any(PathInfo::is_ip)
+                {
+                    tracing::info!("break");
+                    return;
+                }
+            }
+        }
+
+        // Verify that both connections are notified of path changes and get an IP and a relay path.
+        tokio::join!(
+            async {
+                tokio::time::timeout(Duration::from_secs(1), wait_for_paths(&mut paths_server))
+                    .instrument(error_span!("paths-server"))
+                    .await
+                    .unwrap()
+            },
+            async {
+                tokio::time::timeout(Duration::from_secs(1), wait_for_paths(&mut paths_client))
+                    .instrument(error_span!("paths-client"))
+                    .await
+                    .unwrap()
+            }
+        );
+
+        // Close the client connection.
+        tracing::info!("close client conn");
+        conn_client.close(0u32.into(), b"");
+
+        // Verify that the path watch streams close.
+        assert_eq!(
+            tokio::time::timeout(Duration::from_secs(1), paths_client.next())
+                .await
+                .unwrap(),
+            None
+        );
+        assert_eq!(
+            tokio::time::timeout(Duration::from_secs(1), paths_server.next())
+                .await
+                .unwrap(),
+            None
+        );
+
+        server.close().await;
+        client.close().await;
+
+        Ok(())
+    }
 }

iroh/src/magicsock.rs

@@ -64,7 +64,7 @@ use crate::{
     disco::{self, SendAddr},
     discovery::{ConcurrentDiscovery, Discovery, EndpointData, UserData},
     key::{DecryptionError, SharedSecret, public_ed_box, secret_ed_box},
-    magicsock::endpoint_map::PathsWatchable,
+    magicsock::endpoint_map::PathsWatcher,
     metrics::EndpointMetrics,
     net_report::{self, IfStateDetails, Report},
 };
@@ -272,14 +272,14 @@ impl MagicSock {
     /// The actor is responsible for holepunching and opening additional paths to this
     /// connection.
     ///
-    /// Returns a future that resolves to [`PathsWatchable`].
+    /// Returns a future that resolves to [`PathsWatcher`].
     ///
     /// The returned future is `'static`, so it can be stored without being liftetime-bound to `&self`.
     pub(crate) fn register_connection(
         &self,
         remote: EndpointId,
         conn: WeakConnectionHandle,
-    ) -> impl Future<Output = Result<PathsWatchable, EndpointStateActorStoppedError>> + Send + 'static
+    ) -> impl Future<Output = Result<PathsWatcher, EndpointStateActorStoppedError>> + Send + 'static
     {
         let (tx, rx) = oneshot::channel();
         let sender = self.endpoint_map.endpoint_state_actor(remote);

iroh/src/magicsock/endpoint_map.rs

@@ -11,23 +11,22 @@ use serde::{Deserialize, Serialize};
 use tokio::sync::mpsc;
 use tracing::warn;
 
+// #[cfg(any(test, feature = "test-utils"))]
+// use crate::endpoint::PathSelection;
+pub(super) use self::endpoint_state::EndpointStateMessage;
+pub(crate) use self::endpoint_state::PathsWatcher;
+use self::endpoint_state::{EndpointStateActor, EndpointStateHandle};
+pub use self::endpoint_state::{PathInfo, PathInfoList};
 use super::{
     DirectAddr, DiscoState, MagicsockMetrics,
     mapped_addrs::{AddrMap, EndpointIdMappedAddr, RelayMappedAddr},
     transports::{self, TransportsSender},
 };
-use crate::disco::{self};
-// #[cfg(any(test, feature = "test-utils"))]
-// use crate::endpoint::PathSelection;
+use crate::disco;
 
 mod endpoint_state;
 mod path_state;
 
-pub(super) use endpoint_state::EndpointStateMessage;
-pub(crate) use endpoint_state::PathsWatchable;
-use endpoint_state::{EndpointStateActor, EndpointStateHandle};
-pub use endpoint_state::{PathInfo, PathInfoList};
-
 // TODO: use this
 // /// Number of endpoints that are inactive for which we keep info about. This limit is enforced
 // /// periodically via [`NodeMap::prune_inactive`].

iroh/src/magicsock/endpoint_map/endpoint_state.rs

@@ -335,7 +335,7 @@ impl EndpointStateActor {
     async fn handle_msg_add_connection(
         &mut self,
         handle: WeakConnectionHandle,
-        tx: oneshot::Sender<PathsWatchable>,
+        tx: oneshot::Sender<PathsWatcher>,
     ) {
         let pub_open_paths = Watchable::default();
         if let Some(conn) = handle.upgrade() {
@@ -397,10 +397,11 @@ impl EndpointStateActor {
             }
             self.trigger_holepunching().await;
         }
-        tx.send(PathsWatchable {
-            open_paths: pub_open_paths,
-            selected_path: self.selected_path.clone(),
-        })
+        tx.send(PathsWatcher::new(
+            pub_open_paths.watch(),
+            self.selected_path.watch(),
+            handle,
+        ))
         .ok();
     }
 
@@ -1002,7 +1003,7 @@ pub(crate) enum EndpointStateMessage {
     /// needed, any new paths discovered via holepunching will be added.  And closed paths
     /// will be removed etc.
     #[debug("AddConnection(..)")]
-    AddConnection(WeakConnectionHandle, oneshot::Sender<PathsWatchable>),
+    AddConnection(WeakConnectionHandle, oneshot::Sender<PathsWatcher>),
     /// Adds a [`EndpointAddr`] with locations where the endpoint might be reachable.
     AddEndpointAddr(EndpointAddr, Source),
     /// Process a received DISCO CallMeMaybe message.
@@ -1128,38 +1129,70 @@ impl ConnectionState {
     }
 }
 
-/// Watchables for the open paths and selected transmission path in a connection.
+/// Watcher for the open paths and selected transmission path in a connection.
 ///
 /// This is stored in the [`Connection`], and the watchables are set from within the endpoint state actor.
 ///
+/// Internally, this contains a boxed-mapped-joined watcher over the open paths in the connection and the
+/// selected path to the remote endpoint. The watcher is boxed because the mapped-joined watcher with
+/// `SmallVec<PathInfoList>` has a size of over 800 bytes, which we don't want to put upon the [`Connection`].
+///
 /// [`Connection`]: crate::endpoint::Connection
-#[derive(Debug, Default, Clone)]
-pub(crate) struct PathsWatchable {
-    /// Watchable for the open paths (in this connection).
-    open_paths: Watchable<PathAddrList>,
-    /// Watchable for the selected transmission path (global for this remote endpoint).
-    selected_path: Watchable<Option<transports::Addr>>,
+#[derive(Clone, derive_more::Debug)]
+#[debug("PathsWatcher")]
+#[allow(clippy::type_complexity)]
+pub(crate) struct PathsWatcher(
+    Box<
+        n0_watcher::Map<
+            (
+                n0_watcher::Direct<PathAddrList>,
+                n0_watcher::Direct<Option<transports::Addr>>,
+            ),
+            PathInfoList,
+        >,
+    >,
+);
+
+impl n0_watcher::Watcher for PathsWatcher {
+    type Value = PathInfoList;
+
+    fn get(&mut self) -> Self::Value {
+        self.0.get()
+    }
+
+    fn is_connected(&self) -> bool {
+        self.0.is_connected()
+    }
+
+    fn poll_updated(
+        &mut self,
+        cx: &mut std::task::Context<'_>,
+    ) -> Poll<Result<Self::Value, n0_watcher::Disconnected>> {
+        self.0.poll_updated(cx)
+    }
 }
 
-impl PathsWatchable {
-    pub(crate) fn watch(
-        &self,
+impl PathsWatcher {
+    fn new(
+        open_paths: n0_watcher::Direct<PathAddrList>,
+        selected_path: n0_watcher::Direct<Option<transports::Addr>>,
         conn_handle: WeakConnectionHandle,
-    ) -> impl Watcher<Value = PathInfoList> + Unpin + Send + Sync + 'static {
-        let joined_watcher = (self.open_paths.watch(), self.selected_path.watch());
-        joined_watcher.map(move |(open_paths, selected_path)| {
-            let selected_path: Option<TransportAddr> = selected_path.map(Into::into);
-            let Some(conn) = conn_handle.upgrade() else {
-                return PathInfoList(Default::default());
-            };
-            let list = open_paths
-                .into_iter()
-                .flat_map(move |(remote, path_id)| {
-                    PathInfo::new(path_id, &conn, remote, selected_path.as_ref())
-                })
-                .collect();
-            PathInfoList(list)
-        })
+    ) -> Self {
+        Self(Box::new(open_paths.or(selected_path).map(
+            move |(open_paths, selected_path)| {
+                let selected_path: Option<TransportAddr> = selected_path.map(Into::into);
+                let Some(conn) = conn_handle.upgrade() else {
+                    return PathInfoList(Default::default());
+                };
+                let list = open_paths
+                    .into_iter()
+                    .flat_map(move |(remote, path_id)| {
+                        PathInfo::new(path_id, &conn, remote, selected_path.as_ref())
+                    })
+                    .collect();
+                PathInfoList(list)
+            },
+        )))
     }
 }