liminal_server/cluster/

sync.rs

//! SRV-005 R5/R6: cross-node subscription propagation and message fan-out.
//!
//! Each channel maps to a beamr process group (the channel name interned into an
//! atom in the default pg scope). A local subscriber joining a channel joins
//! that pg group; beamr broadcasts the membership to connected peers, so every
//! node knows which peers hold subscribers for which channel. Publishing then
//! consults the group's REMOTE members and forwards the published envelope to
//! each one over the existing distribution link.
//!
//! * R5 propagation: [`ClusterSync::on_subscribe`] -> `pg.join`; the broadcast is
//!   beamr's, not ours.
//! * R5 backfill: a fresh `pg.join` only broadcasts on the insert edge, so a peer
//!   that joins the cluster AFTER our subscribers registered would never learn
//!   them. [`ClusterSync::on_peer_join`] re-sends a pg-join control frame for each
//!   of our local members directly to the newcomer.
//! * R5 delivery: [`ClusterSync::on_publish`] sends the envelope (encoded by
//!   [`liminal::channel::encode_envelope`]) as a beamr `SEND` to each remote
//!   member's pid. On that member's home node the frame lands in the subscriber
//!   process's mailbox, which decodes it back into its inbox.
//! * R6 cleanup: when a peer drops, beamr's connection-down hook calls
//!   `purge_remote_node`, so its remote members vanish from every group with no
//!   liminal code on the path. [`ClusterSync::on_peer_leave`] only logs.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};

use beamr::atom::{Atom, AtomTable};
use beamr::distribution::connection::ConnectionManager;
use beamr::distribution::control::{encode_pg_update_frame, encode_send_frame};
use beamr::distribution::pg::{PgRegistry, PgUpdate, RemoteMember};
use beamr::native::ProcessContext;
use beamr::term::Term;

use crate::cluster::discovery::ClusterResolver;
use liminal::channel::{ClusterObserver, encode_envelope};
use liminal::envelope::Envelope;

/// Propagates channel subscriptions and fans published messages out across the
/// cluster via beamr process groups (SRV-005 R5/R6).
#[derive(Clone)]
pub struct ClusterSync {
    inner: Arc<SyncInner>,
}

struct SyncInner {
    pg: Arc<PgRegistry>,
    atoms: Arc<AtomTable>,
    connections: ConnectionManager,
    /// This node's distribution atom — the node component every locally-joined
    /// member carries on the wire.
    local_node: Atom,
    /// The same resolver the scheduler uses, retained so a future enhancement can
    /// register learned peer addresses; held to keep the discovery/scheduler
    /// resolver identity explicit at this seam.
    _resolver: Arc<ClusterResolver>,
    /// Local subscriptions, keyed by channel group atom, each holding the set of
    /// local subscriber pids. The source of truth for R5 peer-join backfill.
    local: Mutex<HashMap<Atom, Vec<u64>>>,
    /// Told each peer join and leave by name (SRV-012 R7: the console
    /// topology's devices and links). Set once, before the initial view.
    peers: OnceLock<PeerListener>,
}

/// Hears a cluster peer join (`true`) or leave (`false`), by node name.
pub type PeerListener = Arc<dyn Fn(&str, bool) + Send + Sync>;

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

impl ClusterSync {
    /// Builds a sync over the scheduler's pg registry and distribution
    /// connections. `local_node` is this node's distribution atom (interned from
    /// the configured node name in the SAME atom table).
    #[must_use]
    pub fn new(
        pg: Arc<PgRegistry>,
        atoms: Arc<AtomTable>,
        connections: ConnectionManager,
        local_node: Atom,
        resolver: Arc<ClusterResolver>,
    ) -> Self {
        Self {
            inner: Arc::new(SyncInner {
                pg,
                atoms,
                connections,
                local_node,
                _resolver: resolver,
                local: Mutex::new(HashMap::new()),
                peers: OnceLock::new(),
            }),
        }
    }

    /// Installs the listener every peer join and leave is reported to. A
    /// second listener is refused (answers false): the first keeps the seat.
    pub fn set_peer_listener(&self, listener: PeerListener) -> bool {
        self.inner.peers.set(listener).is_ok()
    }

    /// Reports peer `name` joining (`present`) or leaving to the listener.
    pub(crate) fn peer_changed(&self, name: &str, present: bool) {
        if let Some(listener) = self.inner.peers.get() {
            listener(name, present);
        }
    }

    /// The default pg scope atom.
    fn scope(&self) -> Atom {
        self.inner.pg.default_scope()
    }

    /// Interns a channel name into its pg group atom.
    fn group(&self, channel: &str) -> Atom {
        self.inner.atoms.intern(channel)
    }

    /// Remote members of `channel`'s group (for publish fan-out and tests).
    #[must_use]
    pub fn remote_targets(&self, channel: &str) -> Vec<RemoteMember> {
        let group = self.group(channel);
        self.inner.pg.remote_members(self.scope(), group)
    }

    /// Records a local subscriber pid under `group` for backfill.
    fn record_local(&self, group: Atom, pid: u64) {
        let mut local = self.lock_local();
        let pids = local.entry(group).or_default();
        if !pids.contains(&pid) {
            pids.push(pid);
        }
        drop(local);
    }

    /// Forgets a local subscriber pid under `group`.
    fn forget_local(&self, group: Atom, pid: u64) {
        let mut local = self.lock_local();
        if let Some(pids) = local.get_mut(&group) {
            pids.retain(|candidate| *candidate != pid);
            if pids.is_empty() {
                local.remove(&group);
            }
        }
    }

    /// A snapshot of every local `(group, pid)` membership for backfill.
    fn local_memberships(&self) -> Vec<(Atom, u64)> {
        let local = self.lock_local();
        local
            .iter()
            .flat_map(|(group, pids)| pids.iter().map(move |pid| (*group, *pid)))
            .collect()
    }

    /// Sends an encoded envelope to one remote member's pid as a beamr `SEND`.
    fn send_to_member(&self, member: RemoteMember, frame_bytes: &[u8]) {
        // The SEND control targets the member's pid_number; on the member's home
        // node that maps to the local subscriber process, whose mailbox receives
        // the binary payload.
        let Some(to_pid) = Term::try_pid(member.pid_number) else {
            tracing::warn!(
                pid_number = member.pid_number,
                "remote member pid out of immediate range; skipping cross-node delivery"
            );
            return;
        };
        let mut context = ProcessContext::new();
        let Ok(payload) = context.alloc_binary(frame_bytes) else {
            tracing::warn!("failed to allocate cross-node envelope payload");
            return;
        };
        let Ok(frame) = encode_send_frame(
            Term::atom(beamr::atom::Atom::OK),
            to_pid,
            payload,
            &self.inner.atoms,
        ) else {
            tracing::warn!("failed to encode cross-node send frame");
            return;
        };
        self.write_frame(member.node, &frame);
    }

    /// Writes a pre-encoded distribution frame to `node`'s connection, if live.
    fn write_frame(&self, node: Atom, frame: &[u8]) {
        let Some(connection) = self.inner.connections.get_connection(node) else {
            // No live link, so the frame is dropped here and nothing above is
            // told. Both causes are accepted. A peer that DEPARTED between the
            // membership snapshot and this send is owed nothing — the frame was
            // addressed to a node that is gone. A STALE membership entry names a
            // node that never had a link, and that delivery is genuinely lost:
            // no retry is possible against an entry that resolves to no
            // connection. The loss window is bounded by R6 cleanup, which drops
            // the stale entry shortly, after which the node stops being selected
            // at all.
            return;
        };
        write_raw_blocking(&connection, frame);
    }

    /// Re-sends a pg-join control frame for one local member to a single node
    /// (R5 backfill). Mirrors the frame beamr broadcasts on a fresh `pg.join`,
    /// but targeted at the newcomer only.
    fn backfill_member(&self, node: Atom, group: Atom, pid: u64) {
        let update = PgUpdate::Join {
            scope: self.scope(),
            group,
            pid,
        };
        if let Ok(frame) = encode_pg_update_frame(update, self.inner.local_node, &self.inner.atoms)
        {
            self.write_frame(node, &frame);
        } else {
            tracing::warn!("failed to encode cluster backfill frame");
        }
    }

    fn lock_local(&self) -> std::sync::MutexGuard<'_, HashMap<Atom, Vec<u64>>> {
        self.inner
            .local
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
    }
}

impl ClusterObserver for ClusterSync {
    fn on_subscribe(&self, channel: &str, subscriber_pid: u64) {
        let group = self.group(channel);
        // pg.join broadcasts the membership to connected peers (R5 propagation).
        self.inner.pg.join(self.scope(), group, subscriber_pid);
        self.record_local(group, subscriber_pid);
        tracing::debug!(
            channel = %channel,
            pid = subscriber_pid,
            "advertised local subscription to cluster"
        );
    }

    fn on_unsubscribe(&self, channel: &str, subscriber_pid: u64) {
        let group = self.group(channel);
        self.inner.pg.leave(self.scope(), group, subscriber_pid);
        self.forget_local(group, subscriber_pid);
        tracing::debug!(
            channel = %channel,
            pid = subscriber_pid,
            "withdrew local subscription from cluster"
        );
    }

    fn on_publish(&self, channel: &str, envelope: &Envelope) {
        // The console's own records are this device's alone (SRV-012 R4): a
        // peer is never sent one, so nothing is encoded for them.
        if !crate::console::tappable(channel) {
            return;
        }
        let targets = self.remote_targets(channel);
        if targets.is_empty() {
            return;
        }
        let frame_bytes = encode_envelope(envelope);
        for member in targets {
            self.send_to_member(member, &frame_bytes);
        }
    }
}

impl ClusterSync {
    /// R5 backfill: re-advertise every local subscription to a newly-joined peer.
    pub fn on_peer_join(&self, node: Atom) {
        for (group, pid) in self.local_memberships() {
            self.backfill_member(node, group, pid);
        }
    }

    /// R6 is automatic (beamr purges the departed node's remote members via its
    /// connection-down hook). This logs the cleanup for operators.
    pub fn on_peer_leave(&self, node: Atom) {
        let name = self
            .inner
            .atoms
            .resolve(node)
            .map_or_else(|| format!("<atom {node:?}>"), str::to_owned);
        tracing::info!(
            peer = %name,
            "peer departed; its remote subscriptions were purged by beamr"
        );
    }
}

/// Writes `frame` to `connection`, driving the async write to completion off any
/// ambient runtime. The membership poll thread that calls this owns no tokio
/// runtime, so a fresh current-thread runtime drives the single write — the same
/// shape as beamr's own synchronous distribution-send bridge.
fn write_raw_blocking(
    connection: &Arc<beamr::distribution::connection::DistConnection>,
    frame: &[u8],
) {
    let connection = Arc::clone(connection);
    let frame = frame.to_vec();
    let write = async move {
        // A lost distribution write is a silently dropped cross-node fan-out
        // or R5 backfill; the connection teardown side effects are write_raw's
        // own, but the loss itself must be visible here.
        if let Err(error) = connection.write_raw(&frame).await {
            tracing::warn!(%error, "cross-node distribution write failed");
        }
    };
    if let Ok(handle) = tokio::runtime::Handle::try_current() {
        if matches!(
            handle.runtime_flavor(),
            tokio::runtime::RuntimeFlavor::MultiThread
        ) {
            tokio::task::block_in_place(|| handle.block_on(write));
            return;
        }
    }
    match tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
    {
        Ok(runtime) => runtime.block_on(write),
        Err(error) => tracing::warn!(error = %error, "failed to build cluster send runtime"),
    }
}

#[cfg(test)]
#[allow(clippy::expect_used, clippy::unwrap_used, clippy::panic)]
mod tests {
    use super::ClusterSync;
    use crate::cluster::discovery::ClusterResolver;
    use beamr::atom::AtomTable;
    use beamr::distribution::connection::ConnectionManager;
    use beamr::distribution::pg::PgRegistry;
    use beamr::distribution::resolver::StaticResolver;
    use liminal::channel::ClusterObserver;
    use std::collections::HashMap;
    use std::sync::Arc;

    fn sync_fixture() -> (ClusterSync, Arc<PgRegistry>, Arc<AtomTable>) {
        let atoms = Arc::new(AtomTable::with_common_atoms());
        let pg = Arc::new(PgRegistry::new(&atoms));
        let connections = ConnectionManager::new(
            Arc::clone(&atoms),
            Arc::new(StaticResolver::new(HashMap::new())),
            "test-cookie",
            "local@127.0.0.1",
            1,
        );
        let local_node = atoms.intern("local@127.0.0.1");
        let resolver = Arc::new(ClusterResolver::new());
        let sync = ClusterSync::new(
            Arc::clone(&pg),
            Arc::clone(&atoms),
            connections,
            local_node,
            resolver,
        );
        (sync, pg, atoms)
    }

    #[test]
    fn subscribe_joins_the_channel_pg_group() {
        let (sync, pg, atoms) = sync_fixture();
        sync.on_subscribe("orders", 42);
        let group = atoms.intern("orders");
        assert_eq!(pg.local_members(pg.default_scope(), group), vec![42]);
    }

    #[test]
    fn unsubscribe_leaves_the_channel_pg_group() {
        let (sync, pg, atoms) = sync_fixture();
        sync.on_subscribe("orders", 42);
        sync.on_unsubscribe("orders", 42);
        let group = atoms.intern("orders");
        assert!(pg.local_members(pg.default_scope(), group).is_empty());
    }

    #[test]
    fn local_memberships_track_subscriptions_for_backfill() {
        let (sync, _pg, _atoms) = sync_fixture();
        sync.on_subscribe("orders", 1);
        sync.on_subscribe("orders", 2);
        sync.on_subscribe("events", 3);
        let mut memberships = sync.local_memberships();
        memberships.sort_by_key(|(group, pid)| (*group, *pid));
        assert_eq!(memberships.len(), 3);
        // After unsubscribing the only member of a group, the group is dropped.
        sync.on_unsubscribe("events", 3);
        let remaining = sync.local_memberships();
        assert_eq!(remaining.len(), 2);
        assert!(remaining.iter().all(|(_, pid)| *pid == 1 || *pid == 2));
    }

    #[test]
    fn remote_targets_empty_without_remote_members() {
        let (sync, _pg, _atoms) = sync_fixture();
        sync.on_subscribe("orders", 1);
        assert!(sync.remote_targets("orders").is_empty());
    }

    /// A failed cross-node distribution write must be warned about, never
    /// silently discarded (message fan-out and R5 backfill would be lost with
    /// no trace). Drives the REAL write path: a live loopback link (the
    /// membership tests' `LivePair` pattern), whose write half is then
    /// deterministically taken down via the public `mark_down_write_timeout`
    /// seam so `write_raw` yields `NotConnected`.
    #[test]
    fn distribution_write_failure_warns() {
        const COOKIE: &str = "sync-write-cookie";
        const PEER_NAME: &str = "sync-peer@127.0.0.1";
        let runtime = tokio::runtime::Builder::new_multi_thread()
            .worker_threads(2)
            .enable_all()
            .build()
            .expect("build loopback test runtime");

        let peer_atoms = Arc::new(AtomTable::with_common_atoms());
        let peer = ConnectionManager::new(
            peer_atoms,
            Arc::new(StaticResolver::new(HashMap::new())),
            COOKIE,
            PEER_NAME,
            1,
        );
        peer.set_runtime_handle(runtime.handle().clone());
        // Kept alive for the whole test: dropping the accept handle would tear
        // the loopback link down underneath the write below.
        let accept = runtime
            .block_on(peer.listen("127.0.0.1:0".parse().expect("loopback addr")))
            .expect("peer binds a loopback listener");

        let mut routes = HashMap::new();
        routes.insert(PEER_NAME.to_owned(), accept.local_addr());
        let dialer_atoms = Arc::new(AtomTable::with_common_atoms());
        let dialer = ConnectionManager::new(
            Arc::clone(&dialer_atoms),
            Arc::new(StaticResolver::new(routes)),
            COOKIE,
            "sync-dialer@127.0.0.1",
            1,
        );
        dialer.set_runtime_handle(runtime.handle().clone());
        runtime
            .block_on(dialer.connect(PEER_NAME))
            .expect("loopback handshake succeeds");

        let peer_node = dialer_atoms.intern(PEER_NAME);
        let connection = dialer
            .get_connection(peer_node)
            .expect("dialer holds the live connection");
        // Take the write half down; the next write_raw returns NotConnected.
        connection.mark_down_write_timeout();

        let log = crate::test_log::CapturedLog::default();
        log.capture(|| super::write_raw_blocking(&connection, b"frame-bytes"));

        let output = log.contents();
        assert!(
            output.contains("WARN") && output.contains("cross-node distribution write failed"),
            "expected the discarded write error to be warned about, got: {output}"
        );
    }

    /// SRV-012 R4: a console publish reaches no peer. Both channels have a
    /// remote member whose pid is out of the immediate range, so a fan-out
    /// that got as far as the member warns; the ordinary channel's does (the
    /// positive control) and the console channel's never starts.
    #[test]
    fn console_publishes_never_reach_the_fan_out() {
        let (sync, pg, atoms) = sync_fixture();
        let remote_node = atoms.intern("node-b@127.0.0.1");
        for channel in ["orders", "console.topology"] {
            let group = atoms.intern(channel);
            pg.apply_remote_join(pg.default_scope(), group, remote_node, u64::MAX, 0);
            assert_eq!(sync.remote_targets(channel).len(), 1);
        }
        let envelope = || {
            liminal::envelope::Envelope::new(
                b"{}".to_vec(),
                None,
                liminal::channel::SchemaId::new(),
                liminal::envelope::PublisherId::default(),
            )
        };
        let skipped = "remote member pid out of immediate range";

        let console = crate::test_log::CapturedLog::default();
        console.capture(|| sync.on_publish("console.topology", &envelope()));
        assert!(
            !console.contents().contains(skipped),
            "a console record fanned out"
        );

        let ordinary = crate::test_log::CapturedLog::default();
        ordinary.capture(|| sync.on_publish("orders", &envelope()));
        assert!(ordinary.contents().contains(skipped), "positive control");
    }

    #[test]
    fn remote_targets_reflect_applied_remote_joins() {
        let (sync, pg, atoms) = sync_fixture();
        let group = atoms.intern("orders");
        let remote_node = atoms.intern("node-b@127.0.0.1");
        pg.apply_remote_join(pg.default_scope(), group, remote_node, 99, 0);
        let targets = sync.remote_targets("orders");
        assert_eq!(targets.len(), 1);
        assert_eq!(targets[0].node, remote_node);
        assert_eq!(targets[0].pid_number, 99);
    }
}