unb_server/

connection.rs

use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, Weak};

use tokio::sync::watch;
use unb_client::{EndpointSet, TransportKind};
use unb_core::{NodeIdentity, RetirementReason};
use unb_runtime::Wire;

use crate::connect::EndpointDialer;
use crate::node::Node;

#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
/// A stable, cloneable failure category for initial connect and reconnect.
pub enum ConnectError {
    #[error("no supported endpoint in set")]
    NoSupportedEndpoint,
    #[error("dial failed for {transport:?}: {message}")]
    Dial {
        transport: TransportKind,
        message: String,
    },
    #[error("dial timed out for {transport:?}")]
    DialTimedOut { transport: TransportKind },
    #[error("peer establishment failed: {message}")]
    Establishment { message: String },
    #[error("peer identity mismatch: expected {expected:?}, got {actual:?}")]
    IdentityMismatch {
        expected: String,
        actual: Option<String>,
    },
    #[error("the owning node is shut down")]
    NodeShutdown,
    #[error("the reconnect attempt was cancelled")]
    Cancelled,
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// Why a durable logical peer connection is currently disconnected.
pub enum DisconnectReason {
    ExplicitDisconnect,
    NodeShutdown,
    SessionRetired { reason: RetirementReason },
    ReconnectFailed { error: ConnectError },
}

#[derive(Clone, Debug, PartialEq, Eq)]
/// The current state of a durable logical peer connection.
pub enum ConnectionStatus {
    Connecting,
    Connected,
    Disconnected { reason: DisconnectReason },
}

#[derive(Clone)]
struct SessionBinding {
    session_id: String,
    wire: Arc<Wire>,
}

struct ReconnectAttempt {
    generation: u64,
    cancellation: unb_runtime::CancellationToken,
    completion: watch::Sender<Option<Result<(), ConnectError>>>,
}

pub(crate) struct ReconnectWait {
    attempt: Arc<ReconnectAttempt>,
}

impl ReconnectWait {
    pub(crate) async fn wait(self) -> Result<(), ConnectError> {
        self.attempt.wait().await
    }
}

impl ReconnectAttempt {
    fn complete(&self, result: Result<(), ConnectError>) {
        self.completion.send_if_modified(move |completion| {
            if completion.is_some() {
                false
            } else {
                *completion = Some(result);
                true
            }
        });
    }

    fn wait(&self) -> impl std::future::Future<Output = Result<(), ConnectError>> + Send + 'static {
        let mut completion = self.completion.subscribe();
        async move {
            loop {
                if let Some(result) = completion.borrow().clone() {
                    return result;
                }
                if completion.changed().await.is_err() {
                    return Err(ConnectError::Cancelled);
                }
            }
        }
    }
}

#[derive(Clone)]
/// A cloneable handle to one verified peer identity.
///
/// Initial [`Node::connect`](crate::Node::connect) eagerly establishes the
/// session and synchronizes routes before returning this handle. Transport loss
/// updates the handle but does not start background work. Call [`reconnect`](Self::reconnect)
/// explicitly to perform one ordered endpoint sweep; concurrent calls share the
/// same in-flight attempt and result.
pub struct PeerConnection {
    node: Weak<Node>,
    peer: Arc<str>,
    endpoints: Arc<RwLock<EndpointSet>>,
    dialer: Arc<RwLock<Option<Arc<dyn EndpointDialer>>>>,
    identity: Arc<RwLock<NodeIdentity>>,
    session: Arc<Mutex<Option<SessionBinding>>>,
    status: watch::Sender<ConnectionStatus>,
    generation: Arc<AtomicU64>,
    reconnect_attempt: Arc<Mutex<Option<Arc<ReconnectAttempt>>>>,
}

impl PeerConnection {
    pub(crate) fn new(
        node: Weak<Node>,
        identity: NodeIdentity,
        endpoints: EndpointSet,
        session_id: String,
        wire: Arc<Wire>,
        dialer: Option<Arc<dyn EndpointDialer>>,
    ) -> PeerConnection {
        let (status, _) = watch::channel(ConnectionStatus::Connected);
        PeerConnection {
            node,
            peer: Arc::from(identity.node_id.as_str()),
            endpoints: Arc::new(RwLock::new(endpoints)),
            dialer: Arc::new(RwLock::new(dialer)),
            identity: Arc::new(RwLock::new(identity)),
            session: Arc::new(Mutex::new(Some(SessionBinding { session_id, wire }))),
            status,
            generation: Arc::new(AtomicU64::new(0)),
            reconnect_attempt: Arc::new(Mutex::new(None)),
        }
    }

    pub fn peer(&self) -> &str {
        &self.peer
    }

    pub fn status(&self) -> ConnectionStatus {
        self.status.borrow().clone()
    }

    pub fn changed(&self) -> impl std::future::Future<Output = ConnectionStatus> + Send + 'static {
        let mut receiver = self.status.subscribe();
        async move {
            if receiver.changed().await.is_err() {
                return receiver.borrow().clone();
            }
            let status = receiver.borrow_and_update().clone();
            status
        }
    }

    pub fn disconnect(&self) {
        let (attempt, binding) = {
            let mut active = self
                .reconnect_attempt
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            self.generation.fetch_add(1, Ordering::AcqRel);
            let attempt = active.take();
            self.publish(ConnectionStatus::Disconnected {
                reason: DisconnectReason::ExplicitDisconnect,
            });
            let binding = self
                .session
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .as_ref()
                .cloned();
            (attempt, binding)
        };
        if let Some(attempt) = attempt {
            attempt.cancellation.cancel();
            attempt.complete(Err(ConnectError::Cancelled));
        }
        if let Some(binding) = binding {
            binding.wire.shutdown();
        }
    }

    pub async fn reconnect(&self) -> Result<(), ConnectError> {
        if self.status() == ConnectionStatus::Connected {
            return Ok(());
        }
        let (attempt, leader) = {
            let mut active = self
                .reconnect_attempt
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            if self.status() == ConnectionStatus::Connected {
                return Ok(());
            }
            if let Some(attempt) = active.as_ref() {
                (attempt.clone(), false)
            } else {
                let generation = self.generation.fetch_add(1, Ordering::AcqRel) + 1;
                let (completion, _) = watch::channel(None);
                let attempt = Arc::new(ReconnectAttempt {
                    generation,
                    cancellation: unb_runtime::CancellationToken::new(),
                    completion,
                });
                *active = Some(attempt.clone());
                self.publish(ConnectionStatus::Connecting);
                (attempt, true)
            }
        };
        if leader {
            let connection = self.clone();
            let running = attempt.clone();
            unb_runtime::RuntimeHandle::current().spawn(async move {
                connection.run_reconnect(running).await;
            });
        }
        attempt.wait().await
    }

    pub(crate) fn owner(&self) -> Option<Arc<Node>> {
        self.node.upgrade()
    }

    pub(crate) fn active_reconnect(&self) -> Option<ReconnectWait> {
        self.reconnect_attempt
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .as_ref()
            .cloned()
            .map(|attempt| ReconnectWait { attempt })
    }

    pub(crate) fn endpoints(&self) -> EndpointSet {
        self.endpoints
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .clone()
    }

    pub(crate) fn replace_endpoints(&self, endpoints: EndpointSet) {
        *self
            .endpoints
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = endpoints;
    }

    pub(crate) fn replace_dialer(&self, dialer: Option<Arc<dyn EndpointDialer>>) {
        *self
            .dialer
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = dialer;
    }

    fn dialer(&self) -> Option<Arc<dyn EndpointDialer>> {
        self.dialer
            .read()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .clone()
    }

    pub(crate) fn bind(&self, identity: NodeIdentity, session_id: String, wire: Arc<Wire>) {
        let _active = self
            .reconnect_attempt
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        self.install(identity, session_id, wire);
    }

    fn install(&self, identity: NodeIdentity, session_id: String, wire: Arc<Wire>) {
        *self
            .identity
            .write()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) = identity;
        *self
            .session
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner()) =
            Some(SessionBinding { session_id, wire });
        self.publish(ConnectionStatus::Connected);
    }

    fn current_session(&self) -> Option<SessionBinding> {
        self.session
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .clone()
    }

    pub(crate) fn retire(&self, session_id: &str, reason: RetirementReason) {
        let _active = self
            .reconnect_attempt
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        let removed = {
            let mut session = self
                .session
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            if session
                .as_ref()
                .is_some_and(|binding| binding.session_id == session_id)
            {
                session.take()
            } else {
                None
            }
        };
        if removed.is_some() && self.status() == ConnectionStatus::Connected {
            self.publish(ConnectionStatus::Disconnected {
                reason: DisconnectReason::SessionRetired { reason },
            });
        }
    }

    pub(crate) fn node_shutdown(&self) {
        let (attempt, binding) = {
            let mut active = self
                .reconnect_attempt
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            self.generation.fetch_add(1, Ordering::AcqRel);
            let attempt = active.take();
            self.publish(ConnectionStatus::Disconnected {
                reason: DisconnectReason::NodeShutdown,
            });
            let binding = self
                .session
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner())
                .take();
            (attempt, binding)
        };
        if let Some(attempt) = attempt {
            attempt.cancellation.cancel();
            attempt.complete(Err(ConnectError::NodeShutdown));
        }
        if let Some(binding) = binding {
            binding.wire.shutdown();
        }
    }

    pub(crate) fn publish(&self, status: ConnectionStatus) {
        if *self.status.borrow() != status {
            self.status.send_replace(status);
        }
    }

    async fn run_reconnect(&self, attempt: Arc<ReconnectAttempt>) {
        let result = async {
            let Some(node) = self.owner() else {
                return Err(ConnectError::NodeShutdown);
            };
            if node.cancellation().is_cancelled() {
                return Err(ConnectError::NodeShutdown);
            }
            if let Some(binding) = self.current_session() {
                tokio::select! {
                    biased;
                    () = attempt.cancellation.cancelled() => return Err(ConnectError::Cancelled),
                    () = binding.wire.closed() => {}
                }
                loop {
                    if node.session(&binding.session_id).await.is_none() {
                        break;
                    }
                    tokio::select! {
                        biased;
                        () = attempt.cancellation.cancelled() => return Err(ConnectError::Cancelled),
                        () = tokio::task::yield_now() => {}
                    }
                }
            }
            node.reconnect_peer(self.peer(), &self.endpoints(), self.dialer())
                .await
        }
        .await;
        let completion = {
            let mut active = self
                .reconnect_attempt
                .lock()
                .unwrap_or_else(|poisoned| poisoned.into_inner());
            let current = active
                .as_ref()
                .is_some_and(|current| Arc::ptr_eq(current, &attempt));
            if !current
                || attempt.cancellation.is_cancelled()
                || self.generation.load(Ordering::Acquire) != attempt.generation
            {
                if let Ok(candidate) = result {
                    candidate.candidate_wire.shutdown();
                }
                None
            } else {
                let result = match result {
                    Ok(candidate) => {
                        self.install(
                            candidate.identity,
                            candidate.selected.session_id,
                            candidate.selected.wire,
                        );
                        Ok(())
                    }
                    Err(_error) if self.status() == ConnectionStatus::Connected => Ok(()),
                    Err(error) => {
                        self.publish(ConnectionStatus::Disconnected {
                            reason: DisconnectReason::ReconnectFailed {
                                error: error.clone(),
                            },
                        });
                        Err(error)
                    }
                };
                active.take();
                Some(result)
            }
        };
        if let Some(result) = completion {
            attempt.complete(result);
        }
    }
}