test(iroh): Make endpoint_relay_connect_loop not flaky (#3402)

## Description

`endpoint_relay_connect_loop` was flaky before. This *should* fix this
issue. It was last marked flaky (but for windows only) in #3354.

What seems to have happened:
- The test spawns a relay server
- The test spawns a server iroh Endpoint
- The server Endpoint does a net report with QAD probes
- For some reason the spawned relay server is very slow in responding to
QAD requests (>3s)
- The QAD probes time out, the server Endpoint ends up without a home
relay
- The test spawns a client iroh Endpoint
- The client Endpoint tries to connect for 30s
- The server Endpoint doesn't do another net report for 30s, though, so
never ends up being reachable
- The client Endpoint times out.

To work around this, I'm starting the server endpoint and waiting for it
to have a relay address.
IMO this is reasonable to do in tests.

I've also made the tests use `Connection::close` and
`Connection::closed` properly and removed `SendStream::stopped` and
`RecvStream::read_to_end(0)` calls.

## Notes

There's some other drive-by changes. Sorry about that, but IMO they're
kinda too small for their own PRs:
- `net_report` thought that `Watchable::set` would return `Err` when
there's no more watchers listening, but that's incorrect: It returns
`Err` from `set` when the value set is the same as the currently stored
value.
- I've also made some small cosmetic changes to `net_report`
- I've removed a 1600 bytes allocation from the hot path of receiving
relay items in `ActiveRelayActor`.

## Change checklist
<!-- Remove any that are not relevant. -->
- [x] Self-review.

Author: matheus23

iroh/src/endpoint.rs

@@ -2298,43 +2298,44 @@ mod tests {
         let server_secret_key = SecretKey::generate(rand::thread_rng());
         let server_peer_id = server_secret_key.public();
 
-        let server = {
-            let relay_map = relay_map.clone();
-            tokio::spawn(
-                async move {
-                    let ep = Endpoint::builder()
-                        .secret_key(server_secret_key)
-                        .alpns(vec![TEST_ALPN.to_vec()])
-                        .relay_mode(RelayMode::Custom(relay_map))
-                        .insecure_skip_relay_cert_verify(true)
-                        .bind()
-                        .await?;
-                    info!("accepting connection");
-                    let incoming = ep.accept().await.e()?;
-                    let conn = incoming.await.e()?;
-                    let mut stream = conn.accept_uni().await.e()?;
-                    let mut buf = [0u8; 5];
-                    stream.read_exact(&mut buf).await.e()?;
-                    info!("Accepted 1 stream, received {buf:?}.  Closing now.");
-                    // close the connection
-                    conn.close(7u8.into(), b"bye");
-
-                    let res = conn.accept_uni().await;
-                    assert_eq!(res.unwrap_err(), quinn::ConnectionError::LocallyClosed);
-
-                    let res = stream.read_to_end(10).await;
-                    assert_eq!(
-                        res.unwrap_err(),
-                        quinn::ReadToEndError::Read(quinn::ReadError::ConnectionLost(
-                            quinn::ConnectionError::LocallyClosed
-                        ))
-                    );
-                    info!("server test completed");
-                    Ok::<_, Error>(())
-                }
-                .instrument(info_span!("test-server")),
-            )
-        };
+        // Wait for the endpoint to be started to make sure it's up before clients try to connect
+        let ep = Endpoint::builder()
+            .secret_key(server_secret_key)
+            .alpns(vec![TEST_ALPN.to_vec()])
+            .relay_mode(RelayMode::Custom(relay_map.clone()))
+            .insecure_skip_relay_cert_verify(true)
+            .bind()
+            .await?;
+        // Wait for the endpoint to be reachable via relay
+        ep.home_relay().initialized().await?;
+
+        let server = tokio::spawn(
+            async move {
+                info!("accepting connection");
+                let incoming = ep.accept().await.e()?;
+                let conn = incoming.await.e()?;
+                let mut stream = conn.accept_uni().await.e()?;
+                let mut buf = [0u8; 5];
+                stream.read_exact(&mut buf).await.e()?;
+                info!("Accepted 1 stream, received {buf:?}.  Closing now.");
+                // close the connection
+                conn.close(7u8.into(), b"bye");
+
+                let res = conn.accept_uni().await;
+                assert_eq!(res.unwrap_err(), quinn::ConnectionError::LocallyClosed);
+
+                let res = stream.read_to_end(10).await;
+                assert_eq!(
+                    res.unwrap_err(),
+                    quinn::ReadToEndError::Read(quinn::ReadError::ConnectionLost(
+                        quinn::ConnectionError::LocallyClosed
+                    ))
+                );
+                info!("server test completed");
+                Ok::<_, Error>(())
+            }
+            .instrument(info_span!("test-server")),
+        );
 
         let client = tokio::spawn(
             async move {
@@ -2439,11 +2440,10 @@ mod tests {
         Ok(())
     }
 
-    #[cfg_attr(windows, ignore = "flaky")]
     #[tokio::test]
     #[traced_test]
     async fn endpoint_relay_connect_loop() -> Result {
-        let start = Instant::now();
+        let test_start = Instant::now();
         let n_clients = 5;
         let n_chunks_per_client = 2;
         let chunk_size = 10;
@@ -2452,65 +2452,67 @@ mod tests {
         let server_secret_key = SecretKey::generate(&mut rng);
         let server_node_id = server_secret_key.public();
 
+        // Make sure the server is bound before having clients connect to it:
+        let ep = Endpoint::builder()
+            .insecure_skip_relay_cert_verify(true)
+            .secret_key(server_secret_key)
+            .alpns(vec![TEST_ALPN.to_vec()])
+            .relay_mode(RelayMode::Custom(relay_map.clone()))
+            .bind()
+            .await?;
+        // Also make sure the server has a working relay connection
+        ep.home_relay().initialized().await?;
+
+        info!(time = ?test_start.elapsed(), "test setup done");
+
         // The server accepts the connections of the clients sequentially.
-        let server = {
-            let relay_map = relay_map.clone();
-            tokio::spawn(
-                async move {
-                    let ep = Endpoint::builder()
-                        .insecure_skip_relay_cert_verify(true)
-                        .secret_key(server_secret_key)
-                        .alpns(vec![TEST_ALPN.to_vec()])
-                        .relay_mode(RelayMode::Custom(relay_map))
-                        .bind()
-                        .await?;
-                    let eps = ep.bound_sockets();
-
-                    info!(me = %ep.node_id().fmt_short(), eps = ?eps, "server listening on");
-                    for i in 0..n_clients {
-                        let round_start = Instant::now();
-                        info!("[server] round {i}");
-                        let incoming = ep.accept().await.e()?;
-                        let conn = incoming.await.e()?;
-                        let node_id = conn.remote_node_id()?;
-                        info!(%i, peer = %node_id.fmt_short(), "accepted connection");
-                        let (mut send, mut recv) = conn.accept_bi().await.e()?;
-                        let mut buf = vec![0u8; chunk_size];
-                        for _i in 0..n_chunks_per_client {
-                            recv.read_exact(&mut buf).await.e()?;
-                            send.write_all(&buf).await.e()?;
-                        }
-                        send.finish().e()?;
-                        send.stopped().await.e()?;
-                        recv.read_to_end(0).await.e()?;
-                        info!(%i, peer = %node_id.fmt_short(), "finished");
-                        info!("[server] round {i} done in {:?}", round_start.elapsed());
+        let server = tokio::spawn(
+            async move {
+                let eps = ep.bound_sockets();
+
+                info!(me = %ep.node_id().fmt_short(), eps = ?eps, "server listening on");
+                for i in 0..n_clients {
+                    let round_start = Instant::now();
+                    info!("[server] round {i}");
+                    let incoming = ep.accept().await.e()?;
+                    let conn = incoming.await.e()?;
+                    let node_id = conn.remote_node_id()?;
+                    info!(%i, peer = %node_id.fmt_short(), "accepted connection");
+                    let (mut send, mut recv) = conn.accept_bi().await.e()?;
+                    let mut buf = vec![0u8; chunk_size];
+                    for _i in 0..n_chunks_per_client {
+                        recv.read_exact(&mut buf).await.e()?;
+                        send.write_all(&buf).await.e()?;
                     }
-                    Ok::<_, Error>(())
+                    send.finish().e()?;
+                    conn.closed().await; // we're the last to send data, so we wait for the other side to close
+                    info!(%i, peer = %node_id.fmt_short(), "finished");
+                    info!("[server] round {i} done in {:?}", round_start.elapsed());
                 }
-                .instrument(error_span!("server")),
-            )
-        };
+                Ok::<_, Error>(())
+            }
+            .instrument(error_span!("server")),
+        );
+
+        let start = Instant::now();
 
         for i in 0..n_clients {
             let round_start = Instant::now();
-            info!("[client] round {}", i);
-            let relay_map = relay_map.clone();
+            info!("[client] round {i}");
             let client_secret_key = SecretKey::generate(&mut rng);
-            let relay_url = relay_url.clone();
             async {
                 info!("client binding");
                 let ep = Endpoint::builder()
                     .alpns(vec![TEST_ALPN.to_vec()])
                     .insecure_skip_relay_cert_verify(true)
-                    .relay_mode(RelayMode::Custom(relay_map))
+                    .relay_mode(RelayMode::Custom(relay_map.clone()))
                     .secret_key(client_secret_key)
                     .bind()
                     .await?;
                 let eps = ep.bound_sockets();
 
                 info!(me = %ep.node_id().fmt_short(), eps=?eps, "client bound");
-                let node_addr = NodeAddr::new(server_node_id).with_relay_url(relay_url);
+                let node_addr = NodeAddr::new(server_node_id).with_relay_url(relay_url.clone());
                 info!(to = ?node_addr, "client connecting");
                 let conn = ep.connect(node_addr, TEST_ALPN).await.e()?;
                 info!("client connected");
@@ -2522,9 +2524,8 @@ mod tests {
                     recv.read_exact(&mut buf).await.e()?;
                     assert_eq!(buf, vec![i; chunk_size]);
                 }
-                send.finish().e()?;
-                send.stopped().await.e()?;
-                recv.read_to_end(0).await.e()?;
+                // we're the last to receive data, so we close
+                conn.close(0u32.into(), b"bye!");
                 info!("client finished");
                 ep.close().await;
                 info!("client closed");

iroh/src/magicsock/transports/relay/actor.rs

@@ -615,15 +615,9 @@ impl ActiveRelayActor {
                         break Ok(());
                     };
                     self.reset_inactive_timeout();
-                    // TODO: This allocation is *very* unfortunate.  But so is the
-                    // allocation *inside* of PacketizeIter...
-                    let batch = std::mem::replace(
-                        &mut send_datagrams_buf,
-                        Vec::with_capacity(SEND_DATAGRAM_BATCH_SIZE),
-                    );
                     // TODO(frando): can we avoid the clone here?
                     let metrics = self.metrics.clone();
-                    let packet_iter = batch.into_iter().map(|item| {
+                    let packet_iter = send_datagrams_buf.drain(..).map(|item| {
                         metrics.send_relay.inc_by(item.datagrams.contents.len() as _);
                         Ok(ClientToRelayMsg::Datagrams {
                             dst_node_id: item.remote_node,

iroh/src/net_report.rs

@@ -450,7 +450,7 @@ impl Client {
                             PROBES_TIMEOUT,
                             run_probe_v4(ip_mapped_addrs, relay_node, quic_client, dns_resolver),
                         ))
-                        .instrument(info_span!("QAD IPv6", %relay_url)),
+                        .instrument(info_span!("QAD-IPv4", %relay_url)),
                 );
             }
 
@@ -468,7 +468,7 @@ impl Client {
                             PROBES_TIMEOUT,
                             run_probe_v6(ip_mapped_addrs, relay_node, quic_client, dns_resolver),
                         ))
-                        .instrument(info_span!("QAD IPv6", %relay_url)),
+                        .instrument(info_span!("QAD-IPv6", %relay_url)),
                 );
             }
         }
@@ -700,30 +700,28 @@ async fn run_probe_v4(
     };
 
     let observer = Watchable::new(None);
-    let ob = observer.clone();
     let node = relay_node.url.clone();
-    let conn2 = conn.clone();
-    let handle = task::spawn(async move {
-        loop {
-            let val = *receiver.borrow();
-            // if we've sent to an ipv4 address, but received an observed address
-            // that is ivp6 then the address is an [IPv4-Mapped IPv6 Addresses](https://doc.rust-lang.org/beta/std/net/struct.Ipv6Addr.html#ipv4-mapped-ipv6-addresses)
-            let val = val.map(|val| SocketAddr::new(val.ip().to_canonical(), val.port()));
-            let latency = conn2.rtt();
-            trace!(?val, ?relay_addr, ?latency, "got addr V4");
-            if ob
-                .set(val.map(|addr| QadProbeReport {
-                    node: node.clone(),
-                    addr,
-                    latency,
-                }))
-                .is_err()
-            {
-                // cancel if the observer is gone
-                break;
-            }
-            if receiver.changed().await.is_err() {
-                break;
+    let handle = task::spawn({
+        let conn = conn.clone();
+        let observer = observer.clone();
+        async move {
+            loop {
+                let val = *receiver.borrow();
+                // if we've sent to an ipv4 address, but received an observed address
+                // that is ivp6 then the address is an [IPv4-Mapped IPv6 Addresses](https://doc.rust-lang.org/beta/std/net/struct.Ipv6Addr.html#ipv4-mapped-ipv6-addresses)
+                let val = val.map(|val| SocketAddr::new(val.ip().to_canonical(), val.port()));
+                let latency = conn.rtt();
+                trace!(?val, ?relay_addr, ?latency, "got addr V4");
+                observer
+                    .set(val.map(|addr| QadProbeReport {
+                        node: node.clone(),
+                        addr,
+                        latency,
+                    }))
+                    .ok();
+                if receiver.changed().await.is_err() {
+                    break;
+                }
             }
         }
     });
@@ -769,30 +767,28 @@ async fn run_probe_v6(
     };
 
     let observer = Watchable::new(None);
-    let ob = observer.clone();
     let node = relay_node.url.clone();
-    let conn2 = conn.clone();
-    let handle = task::spawn(async move {
-        loop {
-            let val = *receiver.borrow();
-            // if we've sent to an ipv4 address, but received an observed address
-            // that is ivp6 then the address is an IPv4-Mapped IPv6 Addresses
-            let val = val.map(|val| SocketAddr::new(val.ip().to_canonical(), val.port()));
-            let latency = conn2.rtt();
-            trace!(?val, ?relay_addr, ?latency, "got addr V6");
-            if ob
-                .set(val.map(|addr| QadProbeReport {
-                    node: node.clone(),
-                    addr,
-                    latency,
-                }))
-                .is_err()
-            {
-                // cancel if the observer is gone
-                break;
-            }
-            if receiver.changed().await.is_err() {
-                break;
+    let handle = task::spawn({
+        let observer = observer.clone();
+        let conn = conn.clone();
+        async move {
+            loop {
+                let val = *receiver.borrow();
+                // if we've sent to an ipv4 address, but received an observed address
+                // that is ivp6 then the address is an [IPv4-Mapped IPv6 Addresses](https://doc.rust-lang.org/beta/std/net/struct.Ipv6Addr.html#ipv4-mapped-ipv6-addresses)
+                let val = val.map(|val| SocketAddr::new(val.ip().to_canonical(), val.port()));
+                let latency = conn.rtt();
+                trace!(?val, ?relay_addr, ?latency, "got addr V6");
+                observer
+                    .set(val.map(|addr| QadProbeReport {
+                        node: node.clone(),
+                        addr,
+                        latency,
+                    }))
+                    .ok();
+                if receiver.changed().await.is_err() {
+                    break;
+                }
             }
         }
     });