datum_cluster/

node.rs

use std::{
    collections::BTreeMap,
    net::{IpAddr, SocketAddr},
    sync::{Arc, Mutex},
    time::{Duration, SystemTime},
};

use datum::{Signal, Subscription, SubscriptionOverflow};
use tokio::{
    net::UdpSocket,
    sync::{mpsc, oneshot},
    task::JoinHandle,
};

use crate::{
    ClusterConfig, ClusterError, ClusterResult, ClusterState, DowningProvider, Member, MemberEvent,
    MemberEventKind, MemberState, TimeoutDowning,
    foca_driver::{
        FocaEvent, FocaIdentity, GossipCommand, WireMemberAnnouncement, incarnation_seed,
        run_gossip_task,
    },
};

pub(crate) enum ActorCommand {
    Foca(Box<FocaEvent>),
    Leave {
        reply: oneshot::Sender<ClusterResult<()>>,
    },
    Abort {
        reply: oneshot::Sender<ClusterResult<()>>,
    },
}

/// Running Datum cluster membership node.
pub struct ClusterNode {
    config: ClusterConfig,
    local_addr: SocketAddr,
    advertise_addr: SocketAddr,
    state: Signal<ClusterState>,
    events: Subscription<MemberEvent>,
    actor: mpsc::UnboundedSender<ActorCommand>,
    tasks: Arc<Mutex<Vec<JoinHandle<()>>>>,
}

impl ClusterNode {
    /// Start a cluster node with [`TimeoutDowning`] using
    /// [`ClusterConfig::downing_timeout`].
    pub async fn start(config: ClusterConfig) -> ClusterResult<Self> {
        let downing = Arc::new(TimeoutDowning::new(config.downing_timeout));
        Self::start_with_downing(config, downing).await
    }

    /// Start a cluster node with a custom downing provider.
    pub async fn start_with_downing(
        config: ClusterConfig,
        downing: Arc<dyn DowningProvider>,
    ) -> ClusterResult<Self> {
        validate_config(&config)?;
        let roles = config.normalized_roles();
        let socket = UdpSocket::bind(config.bind_addr).await?;
        let local_addr = socket.local_addr()?;
        let advertise_addr = normalize_advertise_addr(config.advertise_addr, local_addr);
        let incarnation = incarnation_seed();
        let identity = FocaIdentity {
            node_id: config.node_id.trim().to_owned(),
            address: advertise_addr,
            agent_addr: config.agent_addr,
            roles: roles.clone(),
            incarnation,
        };

        let now = SystemTime::now();
        let local_member = Member::new(
            identity.node_id.clone(),
            advertise_addr,
            identity.agent_addr,
            roles,
            MemberState::Joining,
            incarnation,
            now,
        );
        let state = Signal::new(ClusterState::new(local_member.clone()))?;
        let events = Subscription::new(
            MemberEvent::new(0, MemberEventKind::Initialized, local_member.clone(), now),
            config.event_buffer.max(1),
            SubscriptionOverflow::Backpressure,
        )?;

        let (actor_tx, actor_rx) = mpsc::unbounded_channel();
        let (gossip_tx, gossip_rx) = mpsc::channel(128);
        let tasks = Arc::new(Mutex::new(Vec::new()));

        let actor_state = MembershipActor {
            config: config.clone(),
            state: state.clone(),
            events: events.clone(),
            members: ClusterState::new(local_member.clone()).members,
            local_node_id: local_member.node_id.clone(),
            local_incarnation: local_member.incarnation,
            next_sequence: 1,
            downing,
            gossip: gossip_tx,
        };

        let actor_task = tokio::spawn(run_membership_actor(actor_state, actor_rx));
        let gossip_task = tokio::spawn(run_gossip_task(
            socket,
            identity,
            config.clone(),
            actor_tx.clone(),
            gossip_rx,
        ));

        {
            let mut locked = tasks.lock().expect("cluster task list poisoned");
            locked.push(actor_task);
            locked.push(gossip_task);
        }

        Ok(Self {
            config,
            local_addr,
            advertise_addr,
            state,
            events,
            actor: actor_tx,
            tasks,
        })
    }

    /// Local UDP address bound by the node.
    #[must_use]
    pub const fn local_addr(&self) -> SocketAddr {
        self.local_addr
    }

    /// Address advertised to peers.
    #[must_use]
    pub const fn advertise_addr(&self) -> SocketAddr {
        self.advertise_addr
    }

    /// Node id.
    #[must_use]
    pub fn node_id(&self) -> &str {
        &self.config.node_id
    }

    /// Current-state feed.
    #[must_use]
    pub fn state(&self) -> Signal<ClusterState> {
        self.state.clone()
    }

    /// Lossless membership event feed.
    #[must_use]
    pub fn events(&self) -> Subscription<MemberEvent> {
        self.events.clone()
    }

    /// Current immutable membership snapshot.
    #[must_use]
    pub fn current_state(&self) -> Arc<ClusterState> {
        self.state.get()
    }

    /// Gracefully leave the cluster.
    pub async fn leave(&self) -> ClusterResult<()> {
        let (reply, receiver) = oneshot::channel();
        self.actor
            .send(ActorCommand::Leave { reply })
            .map_err(|_| ClusterError::Stopped)?;
        receiver.await.map_err(|_| ClusterError::Stopped)?
    }

    /// Stop this node without sending a graceful leave. This is useful for
    /// crash/failure tests.
    pub async fn abort(&self) -> ClusterResult<()> {
        let (reply, receiver) = oneshot::channel();
        self.actor
            .send(ActorCommand::Abort { reply })
            .map_err(|_| ClusterError::Stopped)?;
        receiver.await.map_err(|_| ClusterError::Stopped)?
    }

    /// Alias for [`ClusterNode::abort`].
    pub async fn shutdown(&self) -> ClusterResult<()> {
        self.abort().await
    }
}

impl Drop for ClusterNode {
    fn drop(&mut self) {
        if let Ok(mut tasks) = self.tasks.lock() {
            for task in tasks.drain(..) {
                task.abort();
            }
        }
    }
}

struct MembershipActor {
    config: ClusterConfig,
    state: Signal<ClusterState>,
    events: Subscription<MemberEvent>,
    members: BTreeMap<String, Member>,
    local_node_id: String,
    local_incarnation: u64,
    next_sequence: u64,
    downing: Arc<dyn DowningProvider>,
    gossip: mpsc::Sender<GossipCommand>,
}

async fn run_membership_actor(
    mut actor: MembershipActor,
    mut receiver: mpsc::UnboundedReceiver<ActorCommand>,
) {
    actor.mark_local_up();
    let mut downing_tick = tokio::time::interval(downing_interval(actor.config.downing_timeout));

    loop {
        tokio::select! {
            _ = downing_tick.tick() => {
                actor.apply_downing();
            }
            Some(command) = receiver.recv() => {
                match command {
                    ActorCommand::Foca(event) => actor.apply_foca_event(*event),
                    ActorCommand::Leave { reply } => {
                        let result = actor.leave().await;
                        let _ = reply.send(result);
                    }
                    ActorCommand::Abort { reply } => {
                        let result = actor.abort().await;
                        let _ = reply.send(result);
                        return;
                    }
                }
            }
            else => return,
        }
    }
}

impl MembershipActor {
    fn mark_local_up(&mut self) {
        self.transition_existing(
            &self.local_node_id.clone(),
            MemberState::Up,
            MemberEventKind::MemberUp,
        );
    }

    async fn leave(&mut self) -> ClusterResult<()> {
        if self
            .members
            .get(&self.local_node_id)
            .map(|member| member.state == MemberState::Removed)
            .unwrap_or(false)
        {
            return Ok(());
        }

        self.transition_existing(
            &self.local_node_id.clone(),
            MemberState::Leaving,
            MemberEventKind::MemberLeaving,
        );

        let (reply, receiver) = oneshot::channel();
        self.gossip
            .send(GossipCommand::Leave { reply })
            .await
            .map_err(|_| ClusterError::Stopped)?;
        let result = receiver.await.map_err(|_| ClusterError::Stopped)?;

        self.transition_existing(
            &self.local_node_id.clone(),
            MemberState::Exiting,
            MemberEventKind::MemberExiting,
        );
        self.transition_existing(
            &self.local_node_id.clone(),
            MemberState::Removed,
            MemberEventKind::MemberRemoved,
        );
        result
    }

    async fn abort(&mut self) -> ClusterResult<()> {
        let (reply, receiver) = oneshot::channel();
        self.gossip
            .send(GossipCommand::Shutdown { reply })
            .await
            .map_err(|_| ClusterError::Stopped)?;
        receiver.await.map_err(|_| ClusterError::Stopped)?
    }

    fn apply_foca_event(&mut self, event: FocaEvent) {
        match event {
            FocaEvent::MemberUp(identity) | FocaEvent::Rejoin(identity) => {
                self.apply_member_up(identity);
            }
            FocaEvent::MemberDown(identity) => self.apply_member_down(identity),
            FocaEvent::Rename(_before, after) => self.apply_member_up(after),
            FocaEvent::ProbeFailed(identity) => self.apply_probe_failed(identity),
            FocaEvent::GracefulState(announcement) => self.apply_graceful_state(announcement),
            FocaEvent::GossipError(_error) => {}
            FocaEvent::Active | FocaEvent::Idle | FocaEvent::Defunct => {}
        }
    }

    fn apply_member_up(&mut self, identity: FocaIdentity) {
        if self.is_local_identity(&identity) {
            return;
        }

        let now = SystemTime::now();
        let Some(existing) = self.members.get(&identity.node_id).cloned() else {
            let joining = member_from_identity(&identity, MemberState::Joining, now);
            self.publish_member(joining, MemberEventKind::MemberJoining);
            let up = member_from_identity(&identity, MemberState::Up, now);
            self.publish_member(up, MemberEventKind::MemberUp);
            return;
        };

        if identity.incarnation < existing.incarnation {
            return;
        }

        let mut member = member_from_identity(&identity, MemberState::Up, now);
        let rejoined = identity.incarnation > existing.incarnation
            && matches!(existing.state, MemberState::Down | MemberState::Removed);
        if rejoined {
            self.publish_member(member, MemberEventKind::MemberRejoined);
            return;
        }

        member.state = if matches!(existing.state, MemberState::Leaving | MemberState::Exiting) {
            existing.state
        } else {
            MemberState::Up
        };
        let changed_metadata = existing.address != member.address
            || existing.agent_addr != member.agent_addr
            || existing.roles != member.roles
            || existing.incarnation != member.incarnation;
        let became_reachable = existing.unreachable;
        if became_reachable {
            member.state = MemberState::Up;
            self.publish_member(member, MemberEventKind::MemberReachable);
        } else if changed_metadata {
            self.publish_member(member, MemberEventKind::MemberUpdated);
        } else if existing.state != MemberState::Up {
            self.publish_member(member, MemberEventKind::MemberUp);
        }
    }

    fn apply_member_down(&mut self, identity: FocaIdentity) {
        if self.is_local_identity(&identity) {
            return;
        }
        let Some(existing) = self.members.get(&identity.node_id).cloned() else {
            return;
        };
        if identity.incarnation < existing.incarnation {
            return;
        }

        match existing.state {
            MemberState::Leaving | MemberState::Exiting => {
                self.transition_existing(
                    &identity.node_id,
                    MemberState::Exiting,
                    MemberEventKind::MemberExiting,
                );
                self.transition_existing(
                    &identity.node_id,
                    MemberState::Removed,
                    MemberEventKind::MemberRemoved,
                );
            }
            MemberState::Removed => {}
            MemberState::Down => {}
            MemberState::Joining | MemberState::Up => {
                let mut member = existing;
                member.state = MemberState::Down;
                member.unreachable = true;
                member.unreachable_since = member.unreachable_since.or(Some(SystemTime::now()));
                member.seen_at = SystemTime::now();
                self.publish_member(member, MemberEventKind::MemberDown);
            }
        }
    }

    fn apply_probe_failed(&mut self, identity: FocaIdentity) {
        if self.is_local_identity(&identity) {
            return;
        }
        let Some(mut member) = self.members.get(&identity.node_id).cloned() else {
            return;
        };
        if identity.incarnation < member.incarnation
            || matches!(member.state, MemberState::Down | MemberState::Removed)
        {
            return;
        }
        if !member.unreachable {
            let now = SystemTime::now();
            member.unreachable = true;
            member.unreachable_since = Some(now);
            member.seen_at = now;
            self.publish_member(member, MemberEventKind::MemberUnreachable);
        }
    }

    fn apply_graceful_state(&mut self, announcement: WireMemberAnnouncement) {
        if announcement.node_id == self.local_node_id
            && announcement.incarnation == self.local_incarnation
        {
            return;
        }
        let now = SystemTime::now();
        let state = announcement.member_state();
        if let Some(existing) = self.members.get(&announcement.node_id)
            && announcement.incarnation < existing.incarnation
        {
            return;
        }
        let mut member = Member::new(
            announcement.node_id,
            announcement.address,
            announcement.agent_addr,
            announcement.roles,
            state,
            announcement.incarnation,
            now,
        );
        member.unreachable = false;
        member.unreachable_since = None;
        self.publish_member(member, MemberEventKind::MemberLeaving);
    }

    fn apply_downing(&mut self) {
        let now = SystemTime::now();
        let to_down = self
            .members
            .values()
            .filter(|member| member.node_id != self.local_node_id)
            .filter(|member| member.unreachable)
            .filter(|member| matches!(member.state, MemberState::Joining | MemberState::Up))
            .filter_map(|member| {
                let since = member.unreachable_since?;
                let unreachable_for = now.duration_since(since).unwrap_or_default();
                self.downing
                    .should_down(member, unreachable_for, now)
                    .then(|| member.node_id.clone())
            })
            .collect::<Vec<_>>();

        for node_id in to_down {
            if let Some(mut member) = self.members.get(&node_id).cloned() {
                member.state = MemberState::Down;
                member.seen_at = now;
                self.publish_member(member, MemberEventKind::MemberDown);
            }
        }
    }

    fn transition_existing(&mut self, node_id: &str, state: MemberState, kind: MemberEventKind) {
        let Some(mut member) = self.members.get(node_id).cloned() else {
            return;
        };
        if member.state == state {
            return;
        }
        member.state = state;
        member.seen_at = SystemTime::now();
        if matches!(
            state,
            MemberState::Up | MemberState::Leaving | MemberState::Removed
        ) {
            member.unreachable = false;
            member.unreachable_since = None;
        }
        self.publish_member(member, kind);
    }

    fn publish_member(&mut self, member: Member, kind: MemberEventKind) {
        let now = SystemTime::now();
        self.members.insert(member.node_id.clone(), member.clone());
        let state = ClusterState {
            self_node: self.local_node_id.clone(),
            members: self.members.clone(),
        };
        let event = MemberEvent::new(self.next_sequence, kind, member, now);
        self.next_sequence = self.next_sequence.wrapping_add(1).max(1);
        let _ = self.state.set(state);
        let _ = self.events.set(event);
    }

    fn is_local_identity(&self, identity: &FocaIdentity) -> bool {
        identity.node_id == self.local_node_id && identity.incarnation == self.local_incarnation
    }
}

fn validate_config(config: &ClusterConfig) -> ClusterResult<()> {
    if config.node_id.trim().is_empty() {
        return Err(ClusterError::InvalidConfig(
            "node_id must not be empty".to_owned(),
        ));
    }
    if config.gossip_interval.is_zero() {
        return Err(ClusterError::InvalidConfig(
            "gossip_interval must be non-zero".to_owned(),
        ));
    }
    if config.probe_timeout.is_zero() || config.probe_timeout >= config.gossip_interval {
        return Err(ClusterError::InvalidConfig(
            "probe_timeout must be non-zero and smaller than gossip_interval".to_owned(),
        ));
    }
    if config.downing_timeout.is_zero() {
        return Err(ClusterError::InvalidConfig(
            "downing_timeout must be non-zero".to_owned(),
        ));
    }
    Ok(())
}

fn normalize_advertise_addr(advertise: SocketAddr, local: SocketAddr) -> SocketAddr {
    let ip = if advertise.ip().is_unspecified() {
        local.ip()
    } else {
        advertise.ip()
    };
    let port = if advertise.port() == 0 {
        local.port()
    } else {
        advertise.port()
    };
    SocketAddr::new(normalize_unspecified_loopback(ip, local.ip()), port)
}

fn normalize_unspecified_loopback(ip: IpAddr, local_ip: IpAddr) -> IpAddr {
    if ip.is_unspecified() { local_ip } else { ip }
}

fn member_from_identity(identity: &FocaIdentity, state: MemberState, now: SystemTime) -> Member {
    Member::new(
        identity.node_id.clone(),
        identity.address,
        identity.agent_addr,
        identity.roles.clone(),
        state,
        identity.incarnation,
        now,
    )
}

fn downing_interval(timeout: Duration) -> Duration {
    (timeout / 4).clamp(Duration::from_millis(10), Duration::from_millis(100))
}