use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap},
fmt,
net::SocketAddr,
num::{NonZeroU8, NonZeroUsize},
time::{SystemTime, UNIX_EPOCH},
};
use bincode::config::standard;
use datum_net::Datagram;
use foca::{
AccumulatingRuntime, BincodeCodec, BroadcastHandler, Config as FocaConfig, Foca, Identity,
Invalidates, OwnedNotification, PeriodicParams, Timer,
};
use rand::{SeedableRng, rngs::SmallRng};
use serde::{Deserialize, Serialize};
use tokio::{
net::UdpSocket,
sync::{mpsc, oneshot},
time::{Instant, sleep_until},
};
use crate::{ClusterConfig, ClusterError, ClusterResult, MemberState};
pub(crate) type ActorSender = mpsc::UnboundedSender<crate::node::ActorCommand>;
type ClusterCodec = BincodeCodec<bincode::config::Configuration>;
#[derive(Debug)]
pub(crate) enum GossipCommand {
Leave {
reply: oneshot::Sender<ClusterResult<()>>,
},
Shutdown {
reply: oneshot::Sender<ClusterResult<()>>,
},
}
#[derive(Debug)]
pub(crate) enum FocaEvent {
MemberUp(FocaIdentity),
MemberDown(FocaIdentity),
Rename(FocaIdentity, FocaIdentity),
ProbeFailed(FocaIdentity),
Active,
Idle,
Defunct,
Rejoin(FocaIdentity),
GracefulState(WireMemberAnnouncement),
GossipError(String),
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct FocaIdentity {
pub(crate) node_id: String,
pub(crate) address: SocketAddr,
pub(crate) agent_addr: Option<SocketAddr>,
pub(crate) roles: Vec<String>,
pub(crate) incarnation: u64,
}
impl FocaIdentity {
pub(crate) fn seed(address: SocketAddr) -> Self {
Self {
node_id: format!("seed-{address}"),
address,
agent_addr: None,
roles: Vec::new(),
incarnation: 0,
}
}
}
impl Identity for FocaIdentity {
type Addr = SocketAddr;
fn renew(&self) -> Option<Self> {
Some(Self {
node_id: self.node_id.clone(),
address: self.address,
agent_addr: self.agent_addr,
roles: self.roles.clone(),
incarnation: self.incarnation.wrapping_add(1),
})
}
fn addr(&self) -> Self::Addr {
self.address
}
fn win_addr_conflict(&self, adversary: &Self) -> bool {
self.incarnation > adversary.incarnation
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub(crate) struct WireMemberAnnouncement {
pub(crate) key: BroadcastKey,
pub(crate) node_id: String,
pub(crate) address: SocketAddr,
pub(crate) agent_addr: Option<SocketAddr>,
pub(crate) roles: Vec<String>,
pub(crate) incarnation: u64,
pub(crate) state: WireMemberState,
}
impl WireMemberAnnouncement {
fn leaving(identity: &FocaIdentity, sequence: u64) -> Self {
Self {
key: BroadcastKey {
node_id: identity.node_id.clone(),
incarnation: identity.incarnation,
sequence,
},
node_id: identity.node_id.clone(),
address: identity.address,
agent_addr: identity.agent_addr,
roles: identity.roles.clone(),
incarnation: identity.incarnation,
state: WireMemberState::Leaving,
}
}
pub(crate) fn member_state(&self) -> MemberState {
match self.state {
WireMemberState::Leaving => MemberState::Leaving,
}
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub(crate) enum WireMemberState {
Leaving,
}
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub(crate) struct BroadcastKey {
node_id: String,
incarnation: u64,
sequence: u64,
}
impl Invalidates for BroadcastKey {
fn invalidates(&self, other: &Self) -> bool {
self.node_id == other.node_id
&& (self.incarnation > other.incarnation
|| (self.incarnation == other.incarnation && self.sequence >= other.sequence))
}
}
#[derive(Debug)]
struct BroadcastError(String);
impl fmt::Display for BroadcastError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(formatter)
}
}
impl std::error::Error for BroadcastError {}
struct ClusterBroadcastHandler {
actor: ActorSender,
seen: HashMap<(String, u64), u64>,
}
impl ClusterBroadcastHandler {
fn new(actor: ActorSender) -> Self {
Self {
actor,
seen: HashMap::new(),
}
}
}
impl BroadcastHandler<FocaIdentity> for ClusterBroadcastHandler {
type Key = BroadcastKey;
type Error = BroadcastError;
fn receive_item(
&mut self,
data: &[u8],
sender: Option<&FocaIdentity>,
) -> Result<Option<Self::Key>, Self::Error> {
let (announcement, read): (WireMemberAnnouncement, usize) =
bincode::serde::decode_from_slice(data, standard())
.map_err(|error| BroadcastError(format!("decode cluster broadcast: {error}")))?;
if read != data.len() {
return Err(BroadcastError(
"trailing cluster broadcast bytes".to_owned(),
));
}
let seen_key = (
announcement.key.node_id.clone(),
announcement.key.incarnation,
);
let is_new = self
.seen
.get(&seen_key)
.map(|seen_sequence| *seen_sequence < announcement.key.sequence)
.unwrap_or(true);
if !is_new {
return Ok(None);
}
self.seen.insert(seen_key, announcement.key.sequence);
if sender.is_some() {
let _ = self.actor.send(crate::node::ActorCommand::Foca(Box::new(
FocaEvent::GracefulState(announcement.clone()),
)));
}
Ok(Some(announcement.key))
}
}
pub(crate) async fn run_gossip_task(
socket: UdpSocket,
identity: FocaIdentity,
config: ClusterConfig,
actor: ActorSender,
mut commands: mpsc::Receiver<GossipCommand>,
) {
let mut driver = match GossipDriver::new(socket, identity, config, actor.clone()) {
Ok(driver) => driver,
Err(error) => {
let _ = actor.send(crate::node::ActorCommand::Foca(Box::new(
FocaEvent::GossipError(error.to_string()),
)));
return;
}
};
if let Err(error) = driver.announce_to_seeds().await {
let _ = actor.send(crate::node::ActorCommand::Foca(Box::new(
FocaEvent::GossipError(error.to_string()),
)));
}
loop {
if let Err(error) = driver.handle_due_timers().await {
driver.report_error(error);
}
if let Some(deadline) = driver.timers.next_deadline() {
tokio::select! {
biased;
Some(command) = commands.recv() => {
if driver.handle_command(command).await {
return;
}
}
_ = sleep_until(deadline) => {}
received = driver.receive_datagram() => {
if let Err(error) = received {
driver.report_error(error);
}
}
}
} else {
tokio::select! {
biased;
Some(command) = commands.recv() => {
if driver.handle_command(command).await {
return;
}
}
received = driver.receive_datagram() => {
if let Err(error) = received {
driver.report_error(error);
}
}
else => return,
}
}
}
}
struct GossipDriver {
socket: UdpSocket,
foca: Foca<FocaIdentity, ClusterCodec, SmallRng, ClusterBroadcastHandler>,
runtime: AccumulatingRuntime<FocaIdentity>,
timers: TimerQueue,
recv_buf: Vec<u8>,
seeds: Vec<FocaIdentity>,
actor: ActorSender,
broadcast_sequence: u64,
}
impl GossipDriver {
fn new(
socket: UdpSocket,
identity: FocaIdentity,
config: ClusterConfig,
actor: ActorSender,
) -> ClusterResult<Self> {
let mut foca_config = foca_config(&config)?;
foca_config.periodic_gossip = Some(PeriodicParams {
frequency: config.gossip_interval,
num_members: NonZeroUsize::new(3).expect("non-zero gossip fanout"),
});
foca_config.periodic_announce = Some(PeriodicParams {
frequency: config.gossip_interval.saturating_mul(4),
num_members: NonZeroUsize::new(1).expect("non-zero announce fanout"),
});
let seed = identity.incarnation ^ nanos_since_epoch();
let rng = SmallRng::seed_from_u64(seed);
let seeds = config
.seed_nodes
.iter()
.copied()
.filter(|seed_addr| *seed_addr != identity.address)
.map(FocaIdentity::seed)
.collect::<Vec<_>>();
let recv_buf = vec![0; foca_config.max_packet_size.get()];
let foca = Foca::with_custom_broadcast(
identity,
foca_config,
rng,
BincodeCodec(standard()),
ClusterBroadcastHandler::new(actor.clone()),
);
Ok(Self {
socket,
foca,
runtime: AccumulatingRuntime::new(),
timers: TimerQueue::new(),
recv_buf,
seeds,
actor,
broadcast_sequence: 0,
})
}
async fn announce_to_seeds(&mut self) -> ClusterResult<()> {
let seeds = self.seeds.clone();
for seed in seeds {
if let Err(error) = self.foca.announce(seed, &mut self.runtime) {
self.report_error(ClusterError::Gossip(error.to_string()));
}
self.drain_runtime().await?;
}
Ok(())
}
async fn receive_datagram(&mut self) -> ClusterResult<()> {
let (len, _remote) = self.socket.recv_from(&mut self.recv_buf).await?;
if let Err(error) = self
.foca
.handle_data(&self.recv_buf[..len], &mut self.runtime)
{
self.report_error(ClusterError::Gossip(error.to_string()));
}
self.drain_runtime().await
}
async fn handle_due_timers(&mut self) -> ClusterResult<()> {
let now = Instant::now();
while let Some((_deadline, timer)) = self.timers.pop_due(now) {
if let Err(error) = self.foca.handle_timer(timer, &mut self.runtime) {
self.report_error(ClusterError::Gossip(error.to_string()));
}
self.drain_runtime().await?;
}
Ok(())
}
async fn handle_command(&mut self, command: GossipCommand) -> bool {
match command {
GossipCommand::Leave { reply } => {
let result = self.leave().await;
let _ = reply.send(result);
true
}
GossipCommand::Shutdown { reply } => {
let _ = reply.send(Ok(()));
true
}
}
}
async fn leave(&mut self) -> ClusterResult<()> {
self.broadcast_sequence = self.broadcast_sequence.wrapping_add(1);
let announcement =
WireMemberAnnouncement::leaving(self.foca.identity(), self.broadcast_sequence);
let data = bincode::serde::encode_to_vec(&announcement, standard())
.map_err(|error| ClusterError::Gossip(format!("encode leave broadcast: {error}")))?;
self.foca
.add_broadcast(&data)
.map_err(|error| ClusterError::Gossip(error.to_string()))?;
self.foca
.broadcast(&mut self.runtime)
.map_err(|error| ClusterError::Gossip(error.to_string()))?;
self.drain_runtime().await?;
self.foca
.leave_cluster(&mut self.runtime)
.map_err(|error| ClusterError::Gossip(error.to_string()))?;
self.drain_runtime().await
}
async fn drain_runtime(&mut self) -> ClusterResult<()> {
while let Some((dst, data)) = self.runtime.to_send() {
let datagram = Datagram::new(data.to_vec(), dst.address);
let expected = datagram.payload().len();
let sent = self
.socket
.send_to(datagram.payload(), datagram.remote())
.await?;
if sent != expected {
return Err(ClusterError::Gossip(format!(
"short UDP send to {}: sent {sent} of {expected}",
dst.address
)));
}
}
let now = Instant::now();
while let Some((delay, timer)) = self.runtime.to_schedule() {
self.timers.enqueue(now + delay, timer);
}
while let Some(notification) = self.runtime.to_notify() {
self.emit_notification(notification);
}
Ok(())
}
fn emit_notification(&self, notification: OwnedNotification<FocaIdentity>) {
let event = match notification {
OwnedNotification::MemberUp(identity) => FocaEvent::MemberUp(identity),
OwnedNotification::MemberDown(identity) => FocaEvent::MemberDown(identity),
OwnedNotification::Rename(before, after) => FocaEvent::Rename(before, after),
OwnedNotification::Active => FocaEvent::Active,
OwnedNotification::Idle => FocaEvent::Idle,
OwnedNotification::Defunct => FocaEvent::Defunct,
OwnedNotification::Rejoin(identity) => FocaEvent::Rejoin(identity),
OwnedNotification::ProbeFailed(_, identity) => FocaEvent::ProbeFailed(identity),
OwnedNotification::DataReceived(_) | OwnedNotification::DataSent(_) => return,
};
let _ = self
.actor
.send(crate::node::ActorCommand::Foca(Box::new(event)));
}
fn report_error(&self, error: ClusterError) {
let _ = self.actor.send(crate::node::ActorCommand::Foca(Box::new(
FocaEvent::GossipError(error.to_string()),
)));
}
}
fn foca_config(config: &ClusterConfig) -> ClusterResult<FocaConfig> {
if config.gossip_interval.is_zero() {
return Err(ClusterError::InvalidConfig(
"gossip_interval must be non-zero".to_owned(),
));
}
if config.probe_timeout.is_zero() {
return Err(ClusterError::InvalidConfig(
"probe_timeout must be non-zero".to_owned(),
));
}
if config.probe_timeout >= config.gossip_interval {
return Err(ClusterError::InvalidConfig(
"probe_timeout must be smaller than gossip_interval".to_owned(),
));
}
if config.suspect_timeout.is_zero() {
return Err(ClusterError::InvalidConfig(
"suspect_timeout must be non-zero".to_owned(),
));
}
let max_packet_size = NonZeroUsize::new(config.max_packet_size).ok_or_else(|| {
ClusterError::InvalidConfig("max_packet_size must be non-zero".to_owned())
})?;
Ok(FocaConfig {
probe_period: config.gossip_interval,
probe_rtt: config.probe_timeout,
num_indirect_probes: NonZeroUsize::new(3).expect("non-zero indirect probes"),
max_transmissions: NonZeroU8::new(10).expect("non-zero transmissions"),
suspect_to_down_after: config.suspect_timeout,
remove_down_after: config.remove_down_after,
max_packet_size,
notify_down_members: true,
periodic_announce: None,
periodic_announce_to_down_members: None,
periodic_gossip: None,
})
}
#[derive(Default)]
struct TimerQueue(BinaryHeap<Reverse<(Instant, Timer<FocaIdentity>)>>);
impl TimerQueue {
fn new() -> Self {
Self::default()
}
fn enqueue(&mut self, deadline: Instant, timer: Timer<FocaIdentity>) {
self.0.push(Reverse((deadline, timer)));
}
fn next_deadline(&self) -> Option<Instant> {
self.0.peek().map(|Reverse((deadline, _timer))| *deadline)
}
fn pop_due(&mut self, now: Instant) -> Option<(Instant, Timer<FocaIdentity>)> {
if self
.0
.peek()
.map(|Reverse((deadline, _timer))| *deadline <= now)
.unwrap_or(false)
{
self.0.pop().map(|Reverse(inner)| inner)
} else {
None
}
}
}
pub(crate) fn incarnation_seed() -> u64 {
static NEXT: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(1);
let process_unique = NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
nanos_since_epoch().wrapping_add(process_unique)
}
fn nanos_since_epoch() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_nanos() as u64)
.unwrap_or(0)
}