Merge c70ad9e into 2f924d9

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)
+            },
+        )))
     }
 }