swactor/crates/swactor-gossip/src/actor.rs

use swactor::actor::{ActorAddress, ActorInterface, Ctx};

use crate::message::{GossipMessage, GossipQueryResponse};
use crate::state::GossipState;
use crate::trace::{
    current_tick, record_event, resolve_name, EventLog, GossipEvent, GossipEventKind,
    NameRegistry, NodeSnapshot, TickCounter,
};

pub struct GossipActor {
    state: GossipState,
    peers: Vec<ActorAddress>,
    event_log: Option<EventLog>,
    tick_counter: Option<TickCounter>,
    name_registry: Option<NameRegistry>,
}

impl GossipActor {
    pub fn new() -> Self {
        Self {
            state: GossipState::new(),
            peers: Vec::new(),
            event_log: None,
            tick_counter: None,
            name_registry: None,
        }
    }

    /// Create a traced actor that records events into the shared log.
    pub fn traced(log: EventLog, tick: TickCounter, names: NameRegistry) -> Self {
        Self {
            state: GossipState::new(),
            peers: Vec::new(),
            event_log: Some(log),
            tick_counter: Some(tick),
            name_registry: Some(names),
        }
    }

    fn pick_random_peer(&self) -> Option<ActorAddress> {
        if self.peers.is_empty() {
            return None;
        }
        let mut buf = [0u8; 8];
        getrandom::getrandom(&mut buf).unwrap();
        let idx = usize::from_ne_bytes(buf) % self.peers.len();
        Some(self.peers[idx])
    }

    fn record(&self, addr: ActorAddress, kind: GossipEventKind) {
        if let (Some(log), Some(tick), Some(names)) =
            (&self.event_log, &self.tick_counter, &self.name_registry)
        {
            let event = GossipEvent {
                tick: current_tick(tick),
                node_name: resolve_name(names, addr),
                node_addr: addr,
                kind,
            };
            record_event(log, event);
        }
    }
}

impl Default for GossipActor {
    fn default() -> Self {
        Self::new()
    }
}

impl ActorInterface for GossipActor {
    type Incoming = GossipMessage;
    type Response = ();

    fn handle(&mut self, ctx: &Ctx, msg: GossipMessage) {
        let self_addr = ctx.self_addr();
        match msg {
            GossipMessage::AddPeer(addr) => {
                if !self.peers.contains(&addr) {
                    self.peers.push(addr);
                    self.record(
                        self_addr,
                        GossipEventKind::PeerAdded {
                            peer_name: self
                                .name_registry
                                .as_ref()
                                .map(|r| resolve_name(r, addr))
                                .unwrap_or_default(),
                        },
                    );
                }
            }
            GossipMessage::RemovePeer(addr) => {
                let before = self.peers.len();
                self.peers.retain(|a| *a != addr);
                if self.peers.len() < before {
                    self.record(
                        self_addr,
                        GossipEventKind::PeerRemoved {
                            peer_name: self
                                .name_registry
                                .as_ref()
                                .map(|r| resolve_name(r, addr))
                                .unwrap_or_default(),
                        },
                    );
                }
            }
            GossipMessage::Set { key, value } => {
                self.state.set(key.clone(), value);
                self.record(self_addr, GossipEventKind::LocalSet { key });
            }
            GossipMessage::DoGossipRound => {
                if let Some(peer) = self.pick_random_peer() {
                    self.record(
                        self_addr,
                        GossipEventKind::GossipRoundStarted {
                            target_name: self
                                .name_registry
                                .as_ref()
                                .map(|r| resolve_name(r, peer))
                                .unwrap_or_default(),
                        },
                    );
                    let _ = ctx.send(
                        peer,
                        GossipMessage::Push {
                            from: self_addr,
                            state: self.state.clone(),
                        },
                    );
                } else {
                    self.record(self_addr, GossipEventKind::GossipRoundNoPeers);
                }
            }
            GossipMessage::Push {
                from,
                state: remote,
            } => {
                let keys_updated = self.state.merge(&remote);
                self.record(
                    self_addr,
                    GossipEventKind::PushReceived {
                        from_name: self
                            .name_registry
                            .as_ref()
                            .map(|r| resolve_name(r, from))
                            .unwrap_or_default(),
                        keys_updated,
                    },
                );
            }
            GossipMessage::TakeSnapshot => {
                self.record(
                    self_addr,
                    GossipEventKind::StateSnapshot {
                        snapshot: NodeSnapshot {
                            entries: self.state.entries().clone(),
                            peer_count: self.peers.len(),
                        },
                    },
                );
            }
            GossipMessage::Query { key, reply_to } => {
                self.record(
                    self_addr,
                    GossipEventKind::QueryReceived { key: key.clone() },
                );
                let entry = self.state.get(&key);
                let resp = GossipQueryResponse {
                    key,
                    value: entry.map(|e| e.value.clone()),
                    version: entry.map(|e| e.version),
                };
                let _ = ctx.send(reply_to, resp);
            }
        }
    }
}