liminal-server 0.25.4

//! Shared end ownership lets a revocation retire keyed transports synchronously.

use super::LinkEnd;
use crate::console::{
    TopologyChange,
    records::{Carrier, Link, LinkCounters, LinkState},
};
use crate::link_handshake::MachineIdentity;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, Weak};

#[cfg(test)]
mod tests;

/// Exact carrier transitions, independent of crossing observations.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, serde::Serialize)]
pub struct LinkLifeCounters {
    /// Successfully keyed links.
    pub links_opened: u64,
    /// Links whose peer carrier disappeared.
    pub links_closed_by_peer_gone: u64,
    /// Links closed by pass revocation.
    pub links_closed_by_revoked: u64,
    /// Links ended by a named refusal.
    pub links_closed_by_refused: u64,
    /// Links superseded by another link.
    pub links_closed_by_replaced: u64,
    /// Links ended by their owner's shutdown.
    pub links_closed_by_shutdown: u64,
}

#[derive(Debug, Default)]
struct Totals {
    opened: AtomicU64,
    peer_gone: AtomicU64,
    revoked: AtomicU64,
    refused: AtomicU64,
    replaced: AtomicU64,
    shutdown: AtomicU64,
}

impl Totals {
    fn increment(counter: &AtomicU64) -> Result<(), String> {
        counter
            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |count| {
                count.checked_add(1)
            })
            .map(|_| ())
            .map_err(|_| "link lifecycle counter overflow".to_owned())
    }

    fn snapshot(&self) -> LinkLifeCounters {
        LinkLifeCounters {
            links_opened: self.opened.load(Ordering::Acquire),
            links_closed_by_peer_gone: self.peer_gone.load(Ordering::Acquire),
            links_closed_by_revoked: self.revoked.load(Ordering::Acquire),
            links_closed_by_refused: self.refused.load(Ordering::Acquire),
            links_closed_by_replaced: self.replaced.load(Ordering::Acquire),
            links_closed_by_shutdown: self.shutdown.load(Ordering::Acquire),
        }
    }
}

type Observer = Arc<dyn Fn(TopologyChange) + Send + Sync>;

#[derive(Default)]
struct State {
    end: Option<LinkEnd>,
    observer: Option<(String, Observer)>,
    notifier: Option<Arc<dyn Fn() + Send + Sync>>,
}

/// One keyed link's first end, shared by its carrier and observations.
#[derive(Clone)]
pub struct LinkLife {
    peer: String,
    totals: Arc<Totals>,
    state: Arc<Mutex<State>>,
}

impl std::fmt::Debug for LinkLife {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LinkLife")
            .field("peer", &self.peer)
            .finish_non_exhaustive()
    }
}

impl LinkLife {
    pub(crate) fn activity(&self) -> Result<(), String> {
        let state = self
            .state
            .lock()
            .map_err(|_| "link lifecycle lock poisoned")?;
        if let Some(notifier) = &state.notifier {
            notifier();
        }
        drop(state);
        Ok(())
    }
    pub(crate) fn notify_with(&self, notifier: Arc<dyn Fn() + Send + Sync>) -> Result<(), String> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| "link lifecycle lock poisoned")?;
        state.notifier = Some(notifier);
        if state.end.is_some()
            && let Some(notifier) = &state.notifier
        {
            notifier();
        }
        drop(state);
        Ok(())
    }
    pub(crate) fn peer(&self) -> &str {
        &self.peer
    }
    pub(crate) fn new(peer: MachineIdentity) -> Self {
        let totals = Totals {
            opened: AtomicU64::new(1),
            ..Totals::default()
        };
        Self::with_totals(hex::encode(peer.ed25519_public()), Arc::new(totals))
    }

    fn with_totals(peer: String, totals: Arc<Totals>) -> Self {
        Self {
            peer,
            totals,
            state: Arc::default(),
        }
    }

    pub(crate) fn observe(&self, local: String, observer: Observer) -> Result<(), String> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| "link lifecycle lock poisoned")?;
        state.observer = Some((local, observer));
        self.emit(&state);
        drop(state);
        Ok(())
    }

    pub(crate) fn end(&self, reason: LinkEnd) -> Result<(), String> {
        let mut state = self
            .state
            .lock()
            .map_err(|_| "link lifecycle lock poisoned")?;
        if state.end.is_some() {
            return Ok(());
        }
        let counter = match &reason {
            LinkEnd::PeerGone => &self.totals.peer_gone,
            LinkEnd::Revoked => &self.totals.revoked,
            LinkEnd::Refused(_) => &self.totals.refused,
            LinkEnd::Replaced => &self.totals.replaced,
            LinkEnd::Shutdown => &self.totals.shutdown,
        };
        Totals::increment(counter)?;
        state.end = Some(reason);
        self.emit(&state);
        if let Some(notifier) = &state.notifier {
            notifier();
        }
        drop(state);
        Ok(())
    }

    pub(crate) fn snapshot(&self) -> LinkLifeCounters {
        self.totals.snapshot()
    }

    pub(crate) fn reason(&self) -> Result<Option<String>, String> {
        let state = self
            .state
            .lock()
            .map_err(|_| "link lifecycle lock poisoned")?;
        Ok(state.end.as_ref().map(ToString::to_string))
    }

    fn emit(&self, state: &State) {
        if let Some((local, observer)) = &state.observer {
            observer(TopologyChange::Link {
                link_id: Link::to_peer(local, &self.peer).id,
                peer: self.peer.clone(),
                carrier: Some(Carrier::Wire),
                state: if state.end.is_some() {
                    LinkState::Down
                } else {
                    LinkState::Up
                },
                reason: state.end.as_ref().map(ToString::to_string),
                counters: LinkCounters {
                    lifecycle: Some(self.snapshot()),
                    ..LinkCounters::default()
                },
            });
        }
    }
}

pub trait CloseLinks: Send + Sync + std::fmt::Debug {
    fn end_links(&self, reason: LinkEnd) -> Result<(), String>;
}

/// Weak registrations do not retain a departed connection's keyed transports.
#[derive(Clone, Debug, Default)]
pub struct LinkRegistry {
    connections: Arc<Mutex<BTreeMap<u64, Weak<dyn CloseLinks>>>>,
    totals: Arc<Mutex<BTreeMap<String, Arc<Totals>>>>,
}

impl LinkRegistry {
    pub(crate) fn register(&self, pid: u64, links: &Arc<dyn CloseLinks>) -> Result<(), String> {
        self.connections
            .lock()
            .map_err(|_| "link registry lock poisoned")?
            .insert(pid, Arc::downgrade(links));
        Ok(())
    }

    pub(crate) fn opened(&self, peer: MachineIdentity) -> Result<LinkLife, String> {
        let peer = hex::encode(peer.ed25519_public());
        let mut by_peer = self
            .totals
            .lock()
            .map_err(|_| "link totals lock poisoned")?;
        let totals = Arc::clone(by_peer.entry(peer.clone()).or_default());
        drop(by_peer);
        Totals::increment(&totals.opened)?;
        Ok(LinkLife::with_totals(peer, totals))
    }

    pub(crate) fn close(&self, pid: u64, reason: LinkEnd) {
        let links = self
            .connections
            .lock()
            .map_err(|_| "link registry lock poisoned".to_owned())
            .map(|connections| connections.get(&pid).and_then(Weak::upgrade));
        let result = links.and_then(|links| links.map_or(Ok(()), |links| links.end_links(reason)));
        if let Err(error) = result {
            tracing::error!(pid, %error, "link closure failed");
        }
    }

    pub(crate) fn forget(&self, pid: u64) {
        match self.connections.lock() {
            Ok(mut connections) => {
                connections.remove(&pid);
            }
            Err(error) => tracing::error!(pid, %error, "link registration removal failed"),
        }
    }

    pub(crate) fn remove(&self, pid: u64) {
        self.close(pid, LinkEnd::PeerGone);
        self.forget(pid);
    }
}