pb-mapper-client 0.10.1

//! Shared relay resolution and bounded setup admission; no established data is owned here.

use crate::network::NetworkHints;
use crate::relay::CarrierId;
use crate::relay::RelayTransport;
use pb_mapper_core::config::ResolvedAddrs;
use std::{
    collections::HashMap,
    net::{SocketAddr, ToSocketAddrs},
    sync::{
        Arc, LazyLock, Mutex as StdMutex, Weak,
        atomic::{AtomicBool, Ordering},
    },
};
use tokio::{
    sync::{Mutex, OwnedSemaphorePermit, Semaphore, watch},
    time::{Duration, Instant},
};
use tokio_util::task::AbortOnDropHandle;

const DNS_TTL: Duration = Duration::from_secs(60);
const DNS_RETRY: Duration = Duration::from_secs(2);
const DNS_WAIT: Duration = Duration::from_secs(2);
pub(crate) const CONTROL_SETUPS: usize = 8;
pub(crate) const DATA_SETUPS: usize = 64;

type Lookup = AbortOnDropHandle<std::io::Result<ResolvedAddrs>>;

struct Cache {
    addresses: Option<ResolvedAddrs>,
    resolved_at: Option<Instant>,
    attempted_at: Option<Instant>,
    generation: u64,
    lookup: Option<Lookup>,
    lookup_generation: u64,
}

struct RelayState {
    name: String,
    #[cfg(feature = "relay-quic")]
    quic: Option<Box<crate::relay::quic::Pool>>,
    fixed: bool,
    cache: Mutex<Cache>,
    refresh: AtomicBool,
    failed: AtomicBool,
    success: watch::Sender<u64>,
    network: Arc<NetworkHints>,
    control_slots: Arc<Semaphore>,
    data_slots: Arc<Semaphore>,
    #[cfg(test)]
    resolver: Option<TestResolver>,
}

#[cfg(test)]
type TestResolver = Arc<
    dyn Fn() -> std::pin::Pin<
            Box<dyn std::future::Future<Output = std::io::Result<ResolvedAddrs>> + Send>,
        > + Send
        + Sync,
>;

#[derive(Clone)]
pub(crate) struct RelayEndpoint(Arc<RelayState>);

impl std::fmt::Debug for RelayEndpoint {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("RelayEndpoint").field(&self.0.name).finish()
    }
}
impl std::fmt::Display for RelayEndpoint {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(&self.0.name)
    }
}

impl RelayEndpoint {
    pub(crate) fn name(&self) -> &str {
        &self.0.name
    }

    pub(crate) fn validate(&self) -> bool {
        if self.0.fixed {
            return true;
        }
        self.0.name.rsplit_once(':').is_some_and(|(host, port)| {
            !host.is_empty()
                && !host
                    .chars()
                    .any(|c| c.is_whitespace() || matches!(c, '/' | '[' | ']' | ':'))
                && port.parse::<u16>().is_ok()
        })
    }
    pub(crate) fn shared(name: &str) -> Self {
        Self::registered(name, RelayTransport::Tcp)
    }

    /// One endpoint per relay name and carrier in this process, for as long
    /// as anything holds it: every client of that relay shares its DNS
    /// cache, its setup budget and, over QUIC, its connections.
    fn registered(name: &str, transport: RelayTransport) -> Self {
        type Registry = HashMap<(String, RelayTransport), Weak<RelayState>>;
        static RELAYS: LazyLock<StdMutex<Registry>> =
            LazyLock::new(|| StdMutex::new(HashMap::new()));
        let mut relays = RELAYS.lock().unwrap_or_else(|e| e.into_inner());
        relays.retain(|_, state| state.strong_count() > 0);
        let key = (name.to_string(), transport);
        if let Some(state) = relays.get(&key).and_then(Weak::upgrade) {
            return Self(state);
        }
        let initial = name.parse::<SocketAddr>().ok().map(ResolvedAddrs::from);
        let relay = Self::new(name.to_string(), initial.is_some(), initial, transport);
        relays.insert(key, Arc::downgrade(&relay.0));
        relay
    }

    pub(crate) fn fixed(addresses: ResolvedAddrs) -> Self {
        Self::new(
            addresses.to_string(),
            true,
            Some(addresses),
            RelayTransport::Tcp,
        )
    }

    fn new(
        name: String,
        fixed: bool,
        addresses: Option<ResolvedAddrs>,
        transport: RelayTransport,
    ) -> Self {
        let (success, _) = watch::channel(0);
        #[cfg(not(feature = "relay-quic"))]
        let _ = transport;
        Self(Arc::new(RelayState {
            name,
            #[cfg(feature = "relay-quic")]
            quic: (transport == RelayTransport::Quic).then(Box::default),
            fixed,
            cache: Mutex::new(Cache {
                resolved_at: addresses.as_ref().map(|_| Instant::now()),
                addresses,
                attempted_at: None,
                generation: 0,
                lookup: None,
                lookup_generation: 0,
            }),
            refresh: AtomicBool::new(false),
            failed: AtomicBool::new(false),
            success,
            network: NetworkHints::shared(),
            control_slots: Arc::new(Semaphore::new(CONTROL_SETUPS)),
            data_slots: Arc::new(Semaphore::new(DATA_SETUPS)),
            #[cfg(test)]
            resolver: None,
        }))
    }

    pub(crate) fn start(&self) {
        self.0.network.start();
    }
    pub(crate) fn notify_network_change(&self) {
        self.0.network.notify();
    }
    pub(crate) fn wake(&self) -> RecoveryWake {
        RecoveryWake {
            network: self.0.network.subscribe(),
            success: self.0.success.subscribe(),
        }
    }
    pub(crate) fn network_generation(&self) -> u64 {
        self.0.network.generation()
    }
    pub(crate) fn data_slots(&self) -> Arc<Semaphore> {
        self.0.data_slots.clone()
    }
    pub(crate) fn active_setups(&self) -> (usize, usize) {
        (
            CONTROL_SETUPS - self.0.control_slots.available_permits(),
            DATA_SETUPS - self.0.data_slots.available_permits(),
        )
    }
    pub(crate) fn dns_age(&self) -> Option<Duration> {
        self.0
            .cache
            .try_lock()
            .ok()
            .and_then(|cache| cache.resolved_at.map(|time| time.elapsed()))
    }

    pub(crate) async fn control_permit(&self) -> OwnedSemaphorePermit {
        // These semaphores are never closed; endpoint owners retain them.
        self.0
            .control_slots
            .clone()
            .acquire_owned()
            .await
            .unwrap_or_else(|_| unreachable!("relay setup semaphore is never closed"))
    }

    pub(crate) fn transport_failed(&self, carrier: Option<CarrierId>) {
        self.0.refresh.store(true, Ordering::Release);
        self.0.failed.store(true, Ordering::Release);
        self.carrier_failed(carrier);
    }

    /// An attempt over `carrier` failed in transport: only that connection
    /// can be to blame, and the pool retires it if the relay has gone silent
    /// on it. An attempt that failed before choosing one — resolving the
    /// relay, discovering or dialling QUIC — leaves the pool alone, there
    /// being nothing in it to blame.
    pub(crate) fn carrier_failed(&self, carrier: Option<CarrierId>) {
        let Some(id) = carrier else { return };
        #[cfg(feature = "relay-quic")]
        if let Some(pool) = &self.0.quic {
            pool.carrier_failed(id);
        }
        // Without QUIC there is no carrier, and `CarrierId` has no values.
        #[cfg(not(feature = "relay-quic"))]
        match id {}
    }

    pub(crate) fn protocol_succeeded(&self) {
        if self.0.failed.swap(false, Ordering::AcqRel) {
            self.0
                .success
                .send_modify(|generation| *generation = generation.wrapping_add(1));
        }
    }

    pub(crate) async fn addresses(&self) -> std::io::Result<ResolvedAddrs> {
        let mut cache = self.0.cache.lock().await;
        let generation = self.network_generation();
        let urgent = cache.generation != generation || self.0.refresh.load(Ordering::Acquire);
        let stale = cache
            .resolved_at
            .is_none_or(|time| time.elapsed() >= DNS_TTL)
            || urgent;
        if let Some(addresses) = &cache.addresses
            && (self.0.fixed || (!stale && cache.lookup.is_none()))
        {
            return Ok(addresses.clone());
        }
        if cache.lookup.is_none() {
            if cache
                .attempted_at
                .is_some_and(|time| time.elapsed() < DNS_RETRY)
            {
                return cache.addresses.clone().ok_or_else(|| {
                    std::io::Error::new(
                        std::io::ErrorKind::NotFound,
                        "relay DNS retry is cooling down",
                    )
                });
            }
            cache.attempted_at = Some(Instant::now());
            cache.lookup_generation = generation;
            self.0.refresh.store(false, Ordering::Release);
            let name = self.0.name.clone();
            #[cfg(test)]
            let resolver = self.0.resolver.clone();
            cache.lookup = Some(AbortOnDropHandle::new(tokio::spawn(async move {
                #[cfg(test)]
                if let Some(resolver) = resolver {
                    return resolver().await;
                }
                // Use the OS resolver so hosts, VPN and TUN split DNS remain authoritative.
                // getaddrinfo cannot be cancelled once running. Keep its permit
                // inside the blocking closure, even if the tunnel is dropped.
                static DNS_SLOTS: LazyLock<Arc<Semaphore>> =
                    LazyLock::new(|| Arc::new(Semaphore::new(8)));
                let permit = DNS_SLOTS
                    .clone()
                    .acquire_owned()
                    .await
                    .map_err(std::io::Error::other)?;
                let addresses = tokio::task::spawn_blocking(move || {
                    let _permit = permit;
                    name.to_socket_addrs().map(|addresses| addresses.collect())
                })
                .await
                .map_err(std::io::Error::other)??;
                ResolvedAddrs::from_candidates(addresses).ok_or_else(|| {
                    std::io::Error::new(
                        std::io::ErrorKind::NotFound,
                        "relay DNS returned no addresses",
                    )
                })
            })));
        }
        if !urgent
            && cache
                .lookup
                .as_ref()
                .is_some_and(|lookup| !lookup.is_finished())
            && let Some(addresses) = &cache.addresses
        {
            // Routine TTL refresh must not delay healthy data setup.
            return Ok(addresses.clone());
        }
        let Some(lookup) = cache.lookup.as_mut() else {
            unreachable!("lookup created above")
        };
        let result = match tokio::time::timeout(DNS_WAIT, lookup).await {
            Ok(result) => {
                cache.lookup.take();
                result
                    .map_err(std::io::Error::other)
                    .and_then(|result| result)
            }
            Err(_) => {
                // Keep the lookup owned across caller cancellation/timeouts. In particular,
                // a blocking OS resolver must not spawn another thread on every retry.
                return cache.addresses.clone().ok_or_else(|| {
                    std::io::Error::new(
                        std::io::ErrorKind::TimedOut,
                        "relay DNS query is still pending",
                    )
                });
            }
        };
        match result {
            Ok(addresses) => {
                cache.addresses = Some(addresses.clone());
                cache.resolved_at = Some(Instant::now());
                cache.generation = cache.lookup_generation;
                Ok(addresses)
            }
            Err(error) => cache.addresses.clone().ok_or(error),
        }
    }

    /// A stream to the relay over this endpoint's carrier, at its current
    /// addresses. See [`Self::connect_addresses`].
    pub(crate) async fn connect(
        &self,
        credential: pb_mapper_core::checksum::Credential,
        carrier: &mut Option<CarrierId>,
    ) -> Result<pb_mapper_protocol::transport::RelayStream, crate::relay::ConnectError> {
        *carrier = None;
        self.connect_addresses(&self.addresses().await?, credential, carrier)
            .await
    }
    /// A stream to the relay at `addresses` over this endpoint's carrier.
    ///
    /// Opening a stream authenticates no request. Every stream's first frame
    /// is the caller's PBM2 request under `credential`, which the relay
    /// authenticates before it dispatches anything, on either carrier.
    /// `credential` matters here only over QUIC, where it chooses the pooled
    /// connection, whose first stream it was bound to when dialled; a TCP
    /// connection is anonymous until that first frame.
    ///
    /// `carrier` is that QUIC connection, set as soon as it is chosen, so
    /// that a caller can charge a failure of this attempt to it with
    /// [`Self::transport_failed`] — one its own deadline cut short included.
    /// It stays `None` until then, and over TCP.
    pub(crate) async fn connect_addresses(
        &self,
        addresses: &ResolvedAddrs,
        credential: pb_mapper_core::checksum::Credential,
        carrier: &mut Option<CarrierId>,
    ) -> Result<pb_mapper_protocol::transport::RelayStream, crate::relay::ConnectError> {
        *carrier = None;
        #[cfg(feature = "relay-quic")]
        if let Some(pool) = &self.0.quic {
            return pool
                .connect(addresses, credential, self.network_generation(), carrier)
                .await
                .map(|stream| pb_mapper_protocol::transport::RelayStream::Quic(Box::new(stream)));
        }
        // Anonymous until the caller's first frame; see above.
        let _ = credential;
        crate::addr::connect_tcp(addresses)
            .await
            .map(Into::into)
            .map_err(Into::into)
    }
    pub(crate) fn transport(&self) -> RelayTransport {
        #[cfg(feature = "relay-quic")]
        if self.0.quic.is_some() {
            return RelayTransport::Quic;
        }
        RelayTransport::Tcp
    }
    /// The process's endpoint for this relay over `transport`. Only for an
    /// endpoint from [`Self::shared`], whose name is what it is shared under.
    pub(crate) fn with_transport(&self, transport: RelayTransport) -> std::io::Result<Self> {
        if transport == self.transport() {
            return Ok(self.clone());
        }
        if cfg!(not(feature = "relay-quic")) && transport == RelayTransport::Quic {
            return Err(std::io::Error::new(
                std::io::ErrorKind::Unsupported,
                "SDK built without relay-quic",
            ));
        }
        Ok(Self::registered(&self.0.name, transport))
    }
    /// The UDP port the relay offers QUIC on, asked over TCP whatever this
    /// endpoint's carrier.
    pub(crate) async fn quic_port(
        &self,
        credential: pb_mapper_core::checksum::Credential,
    ) -> Result<Option<u16>, crate::relay::ConnectError> {
        crate::relay::offered_quic_port(&self.addresses().await?, credential).await
    }
    pub(crate) async fn relay_diagnostics(&self) -> crate::relay::RelayDiagnostics {
        #[cfg(feature = "relay-quic")]
        if let Some(pool) = &self.0.quic {
            return pool.diagnostics();
        }
        crate::relay::RelayDiagnostics::empty(self.transport())
    }
}

pub(crate) struct RecoveryWake {
    network: watch::Receiver<u64>,
    success: watch::Receiver<u64>,
}

impl RecoveryWake {
    pub(crate) fn acknowledge(&mut self) {
        self.network.borrow_and_update();
        self.success.borrow_and_update();
    }
    pub(crate) async fn network_changed(&mut self) {
        let _ = self.network.changed().await;
        self.network.borrow_and_update();
    }
    /// Wait for either signal, and say which it was.
    pub(crate) async fn changed(&mut self) -> Wake {
        let network_woke = tokio::select! {
            biased;
            _ = self.success.changed() => false,
            _ = self.network.changed() => true,
        };
        let wake = {
            let network = self.network.borrow_and_update();
            // Success can win the race while a network change is pending as
            // well: take both together, so the change is not lost.
            if network_woke || network.has_changed() {
                Wake::Network(*network)
            } else {
                Wake::Relay
            }
        };
        self.success.borrow_and_update();
        wake
    }
}

/// Why a [`RecoveryWake`] woke.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum Wake {
    /// The local network changed, to this generation.
    Network(u64),
    /// The relay answered again after failures. That says nothing about
    /// whether a direct path would get through now.
    Relay,
}

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};

    fn resolver_endpoint(resolver: TestResolver) -> RelayEndpoint {
        let mut endpoint =
            RelayEndpoint::new("relay.test:7666".into(), false, None, RelayTransport::Tcp);
        let state = Arc::get_mut(&mut endpoint.0).unwrap();
        state.resolver = Some(resolver);
        state.network = NetworkHints::isolated();
        endpoint
    }

    #[tokio::test(start_paused = true)]
    async fn initial_dns_failure_retries_and_changed_records_replace_candidates() {
        let port = Arc::new(AtomicUsize::new(0));
        let queries = Arc::new(AtomicUsize::new(0));
        let endpoint = resolver_endpoint(Arc::new({
            let port = port.clone();
            let queries = queries.clone();
            move || {
                let port = port.load(Ordering::SeqCst);
                queries.fetch_add(1, Ordering::SeqCst);
                Box::pin(async move {
                    if port == 0 {
                        Err(std::io::ErrorKind::NotFound.into())
                    } else {
                        Ok(SocketAddr::from(([127, 0, 0, 1], port as u16)).into())
                    }
                })
            }
        }));
        assert!(endpoint.addresses().await.is_err());
        assert!(endpoint.addresses().await.is_err());
        assert_eq!(queries.load(Ordering::SeqCst), 1);
        port.store(7001, Ordering::SeqCst);
        tokio::time::advance(DNS_RETRY).await;
        assert_eq!(endpoint.addresses().await.unwrap().primary().port(), 7001);
        port.store(7002, Ordering::SeqCst);
        endpoint.transport_failed(None);
        assert_eq!(endpoint.addresses().await.unwrap().primary().port(), 7001);
        tokio::time::advance(DNS_RETRY).await;
        assert_eq!(endpoint.addresses().await.unwrap().primary().port(), 7002);
        port.store(0, Ordering::SeqCst);
        tokio::time::advance(DNS_TTL).await;
        assert_eq!(
            endpoint.addresses().await.unwrap().primary().port(),
            7002,
            "a failed refresh must retain the last candidates"
        );
        tokio::task::yield_now().await;
        assert_eq!(endpoint.addresses().await.unwrap().primary().port(), 7002);
        assert_eq!(queries.load(Ordering::SeqCst), 4);
    }

    #[tokio::test(start_paused = true)]
    async fn cancelled_dns_waiters_reuse_one_owned_lookup() {
        let release = Arc::new(tokio::sync::Notify::new());
        let queries = Arc::new(AtomicUsize::new(0));
        let endpoint = resolver_endpoint(Arc::new({
            let release = release.clone();
            let queries = queries.clone();
            move || {
                let release = release.clone();
                queries.fetch_add(1, Ordering::SeqCst);
                Box::pin(async move {
                    release.notified().await;
                    Ok(SocketAddr::from(([127, 0, 0, 1], 7666)).into())
                })
            }
        }));
        for _ in 0..100 {
            assert!(
                tokio::time::timeout(Duration::from_millis(10), endpoint.addresses())
                    .await
                    .is_err()
            );
        }
        assert_eq!(queries.load(Ordering::SeqCst), 1);
        release.notify_one();
        assert_eq!(endpoint.addresses().await.unwrap().primary().port(), 7666);
        assert_eq!(queries.load(Ordering::SeqCst), 1);
    }

    #[tokio::test(start_paused = true)]
    async fn dropping_an_endpoint_aborts_its_pending_lookup() {
        struct Active(Arc<AtomicUsize>);
        impl Drop for Active {
            fn drop(&mut self) {
                self.0.fetch_sub(1, Ordering::SeqCst);
            }
        }
        let active = Arc::new(AtomicUsize::new(0));
        let endpoint = resolver_endpoint(Arc::new({
            let active = active.clone();
            move || {
                let active = active.clone();
                Box::pin(async move {
                    active.fetch_add(1, Ordering::SeqCst);
                    let _active = Active(active);
                    std::future::pending().await
                })
            }
        }));
        assert!(
            tokio::time::timeout(Duration::from_millis(10), endpoint.addresses())
                .await
                .is_err()
        );
        assert_eq!(active.load(Ordering::SeqCst), 1);
        drop(endpoint);
        tokio::task::yield_now().await;
        assert_eq!(active.load(Ordering::SeqCst), 0);
    }

    #[tokio::test(start_paused = true)]
    async fn long_outage_wakes_immediately_and_repeated_cancellation_releases_capacity() {
        let endpoint = resolver_endpoint(Arc::new(|| {
            Box::pin(async { Err(std::io::ErrorKind::NotFound.into()) })
        }));
        let mut wake = endpoint.wake();
        // Simulate the observed long outage without changing the host network.
        tokio::time::advance(Duration::from_secs(130)).await;
        endpoint.notify_network_change();
        assert_eq!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .unwrap(),
            Wake::Network(endpoint.network_generation())
        );
        for _ in 0..100 {
            let mut permits = Vec::new();
            for _ in 0..CONTROL_SETUPS {
                permits.push(endpoint.control_permit().await);
            }
            assert!(
                tokio::time::timeout(Duration::from_millis(1), endpoint.control_permit())
                    .await
                    .is_err()
            );
            assert_eq!(endpoint.active_setups().0, CONTROL_SETUPS);
            drop(permits);
            assert_eq!(endpoint.active_setups(), (0, 0));
        }
    }

    #[tokio::test(start_paused = true)]
    async fn recovery_wakes_distinguish_and_coalesce_network_changes() {
        let endpoint = resolver_endpoint(Arc::new(|| {
            Box::pin(async { Err(std::io::ErrorKind::NotFound.into()) })
        }));
        let mut wake = endpoint.wake();
        endpoint.protocol_succeeded();
        assert!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .is_err()
        );
        endpoint.transport_failed(None);
        endpoint.protocol_succeeded();
        assert_eq!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .unwrap(),
            Wake::Relay
        );
        endpoint.protocol_succeeded();
        assert!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .is_err()
        );
        endpoint.transport_failed(None);
        endpoint.protocol_succeeded();
        endpoint.notify_network_change();
        assert_eq!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .unwrap(),
            Wake::Network(endpoint.network_generation())
        );
        assert!(
            tokio::time::timeout(Duration::from_millis(1), wake.changed())
                .await
                .is_err()
        );
    }

    #[cfg(feature = "p2p")]
    #[tokio::test]
    async fn relay_recovery_preserves_direct_cooldown_until_network_changes() {
        use crate::sdk::{Client, ConnectRequest, RegisterOptions, RegisterRequest, Transport};
        use pb_mapper_auth::{AuthConfig, AuthRuntime};
        use pb_mapper_core::checksum::Credential;
        use tokio::{
            io::{AsyncReadExt, AsyncWriteExt},
            net::{TcpListener, TcpStream},
        };
        use tokio_util::sync::CancellationToken;

        tokio::time::timeout(Duration::from_secs(8), async {
            let credential = Credential::Admin(pb_mapper_core::checksum::gen_random_key());
            let state_dir = std::env::temp_dir().join(format!(
                "pb-cooldown-{}-{}",
                std::process::id(),
                rand::random::<u64>()
            ));
            let auth = AuthRuntime::start(
                *credential.key(),
                AuthConfig {
                    state_dir: state_dir.clone(),
                    ..Default::default()
                },
            )
            .await
            .unwrap();
            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let relay_addr = listener.local_addr().unwrap();
            let relay = AbortOnDropHandle::new(tokio::spawn(
                pb_mapper_server::run_server_on_listener_with_direct(
                    listener,
                    CancellationToken::new(),
                    None,
                    false,
                    auth,
                    None,
                ),
            ));
            let backend = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let backend_addr = backend.local_addr().unwrap();
            let echo = AbortOnDropHandle::new(tokio::spawn(async move {
                loop {
                    let (mut stream, _) = backend.accept().await.unwrap();
                    let mut byte = [0; 1];
                    stream.read_exact(&mut byte).await.unwrap();
                    stream.write_all(&byte).await.unwrap();
                }
            }));
            let mut client =
                Client::from_credential(relay_addr.to_string(), credential, false, None);
            let mut endpoint = RelayEndpoint::fixed(relay_addr.into());
            Arc::get_mut(&mut endpoint.0).unwrap().network = NetworkHints::isolated();
            Arc::get_mut(&mut client.inner).unwrap().endpoint = endpoint;
            let registration = client
                .register_with_options(
                    RegisterRequest {
                        key: "cooldown".into(),
                        local_addr: backend_addr.to_string(),
                        transport: Transport::Tcp,
                        codec: true,
                        force_namespace: false,
                    },
                    RegisterOptions {
                        allow_direct: false,
                        ..Default::default()
                    },
                )
                .await
                .unwrap();
            registration
                .wait_ready_timeout(Duration::from_secs(3))
                .await
                .unwrap();
            // The SDK owns its listener and cannot accept a bound socket.
            let reservation = TcpListener::bind("127.0.0.1:0").await.unwrap();
            let local = reservation.local_addr().unwrap();
            drop(reservation);
            let tunnel = client
                .connect(ConnectRequest {
                    key: "cooldown".into(),
                    local_addr: local.to_string(),
                    transport: Transport::Tcp,
                })
                .await
                .unwrap();
            tunnel
                .wait_ready_timeout(Duration::from_secs(3))
                .await
                .unwrap();
            let exchange = async || {
                let mut stream = TcpStream::connect(local).await.unwrap();
                stream.write_all(b"x").await.unwrap();
                let mut response = [0; 1];
                stream.read_exact(&mut response).await.unwrap();
                assert_eq!(&response, b"x");
            };
            exchange().await;
            assert_eq!(tunnel.diagnostics().direct.unwrap().fallbacks, 1);
            exchange().await;
            let before = tunnel.diagnostics();
            assert_eq!(before.direct.unwrap().fallbacks, 1);

            // A shared endpoint's recovery must wake the worker's health probe,
            // without claiming that this tunnel's direct path became reachable.
            let endpoint = &client.inner.endpoint;
            endpoint.transport_failed(None);
            endpoint.protocol_succeeded();
            while tunnel.diagnostics().attempts == before.attempts {
                tokio::task::yield_now().await;
            }
            exchange().await;
            let recovered = tunnel.diagnostics();
            assert_eq!(recovered.network_generation, before.network_generation);
            assert_eq!(recovered.direct.unwrap().fallbacks, 1);

            endpoint.notify_network_change();
            while !tunnel.diagnostics().events.iter().any(|event| {
                matches!(
                    event.kind,
                    crate::diagnostics::TunnelEventKind::NetworkChanged
                )
            }) {
                tokio::task::yield_now().await;
            }
            exchange().await;
            let changed = tunnel.diagnostics();
            assert_ne!(changed.network_generation, before.network_generation);
            assert_eq!(changed.direct.unwrap().fallbacks, 2);
            tunnel.stop().await.unwrap();
            registration.stop().await.unwrap();
            drop(echo);
            drop(relay);
            let _ = std::fs::remove_dir_all(state_dir);
        })
        .await
        .unwrap();
    }
}