use std::collections::VecDeque;
use std::net::SocketAddr;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError, Weak};
use kinship_core::Member;
use tokio::sync::Notify;
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum Event {
MemberJoined(Member),
MemberSuspect(Member),
MemberRecovered(Member),
MemberDead(Member),
MemberLeft(Member),
MemberUpdated {
member: Member,
previous_meta: Vec<u8>,
},
NameConflict {
member: Member,
other_addr: SocketAddr,
},
UserMessage {
sender: Member,
data: Vec<u8>,
broadcast: bool,
},
EventsLost(u64),
}
impl Event {
pub(crate) fn from_core(e: kinship_core::Event) -> Option<Self> {
use kinship_core::Event as E;
Some(match e {
E::MemberJoined(m) => Self::MemberJoined(m),
E::MemberSuspect(m) => Self::MemberSuspect(m),
E::MemberRecovered(m) => Self::MemberRecovered(m),
E::MemberDead(m) => Self::MemberDead(m),
E::MemberLeft(m) => Self::MemberLeft(m),
E::MemberUpdated {
member,
previous_meta,
} => Self::MemberUpdated {
member,
previous_meta,
},
E::NameConflict { member, other_addr } => Self::NameConflict { member, other_addr },
E::UserMessage {
sender,
data,
broadcast,
} => Self::UserMessage {
sender,
data,
broadcast,
},
_ => return None,
})
}
}
fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
m.lock().unwrap_or_else(PoisonError::into_inner)
}
#[derive(Debug)]
pub(crate) struct Hub {
capacity: usize,
inner: Mutex<HubState>,
}
#[derive(Debug, Default)]
struct HubState {
subs: Vec<Weak<Sub>>,
closed: bool,
failed: bool,
}
#[derive(Debug, Default)]
struct Sub {
queue: Mutex<Queue>,
notify: Notify,
}
#[derive(Debug, Default)]
struct Queue {
events: VecDeque<Event>,
lost: u64,
closed: bool,
failed: bool,
}
impl Hub {
pub fn new(capacity: usize) -> Self {
Self {
capacity: capacity.max(1),
inner: Mutex::default(),
}
}
pub fn subscribe(&self) -> Events {
let sub = Arc::new(Sub::default());
let mut hub = lock(&self.inner);
if hub.closed {
let mut q = lock(&sub.queue);
q.closed = true;
q.failed = hub.failed;
} else {
hub.subs.push(Arc::downgrade(&sub));
}
Events { sub }
}
pub fn publish(&self, events: &[Event]) {
if events.is_empty() {
return;
}
let mut hub = lock(&self.inner);
hub.subs.retain(|weak| {
let Some(sub) = weak.upgrade() else {
return false;
};
let mut q = lock(&sub.queue);
for e in events {
if q.events.len() == self.capacity {
q.events.pop_front();
q.lost += 1;
}
q.events.push_back(e.clone());
}
drop(q);
sub.notify.notify_one();
true
});
}
pub fn close(&self, failed: bool) {
let mut hub = lock(&self.inner);
hub.closed = true;
hub.failed = failed;
for sub in hub.subs.drain(..).filter_map(|w| w.upgrade()) {
let mut q = lock(&sub.queue);
q.closed = true;
q.failed = failed;
drop(q);
sub.notify.notify_one();
}
}
}
#[derive(Debug)]
pub struct Events {
sub: Arc<Sub>,
}
impl Events {
pub async fn recv(&mut self) -> Option<Event> {
loop {
match self.next() {
Next::Event(e) => return Some(e),
Next::Closed => return None,
Next::Empty => self.sub.notify.notified().await,
}
}
}
pub fn try_recv(&mut self) -> Option<Event> {
match self.next() {
Next::Event(e) => Some(e),
Next::Closed | Next::Empty => None,
}
}
pub fn is_closed(&self) -> bool {
let q = lock(&self.sub.queue);
q.closed && q.lost == 0 && q.events.is_empty()
}
pub fn failed(&self) -> bool {
lock(&self.sub.queue).failed
}
fn next(&self) -> Next {
let mut q = lock(&self.sub.queue);
if q.lost > 0 {
let lost = std::mem::take(&mut q.lost);
return Next::Event(Event::EventsLost(lost));
}
match q.events.pop_front() {
Some(e) => Next::Event(e),
None if q.closed => Next::Closed,
None => Next::Empty,
}
}
}
enum Next {
Event(Event),
Closed,
Empty,
}
#[cfg(test)]
mod tests {
use super::*;
use kinship_core::State;
fn member(name: &str) -> Member {
Member {
name: name.to_owned(),
addr: SocketAddr::from(([127, 0, 0, 1], 1)),
meta: Vec::new(),
state: State::Alive,
incarnation: 0,
epoch: 0,
}
}
fn joined(i: usize) -> Event {
Event::MemberJoined(member(&format!("n{i}")))
}
#[tokio::test]
async fn subscriptions_start_at_the_call_and_end_on_close() {
let hub = Hub::new(8);
hub.publish(&[joined(0)]);
let mut a = hub.subscribe();
hub.publish(&[joined(1)]);
let mut b = hub.subscribe();
hub.publish(&[joined(2)]);
assert_eq!(a.recv().await, Some(joined(1)));
assert_eq!(a.recv().await, Some(joined(2)));
assert_eq!(b.try_recv(), Some(joined(2)));
assert_eq!(b.try_recv(), None);
assert!(!b.is_closed());
hub.close(false);
hub.publish(&[joined(3)]);
assert_eq!(a.recv().await, None);
assert!(b.is_closed());
assert!(!a.failed() && !b.failed());
assert_eq!(hub.subscribe().recv().await, None);
}
#[tokio::test]
async fn a_failed_node_marks_current_and_later_subscriptions() {
let hub = Hub::new(4);
let mut early = hub.subscribe();
hub.publish(&[joined(0)]);
hub.close(true);
assert_eq!(early.recv().await, Some(joined(0)));
assert_eq!(early.recv().await, None);
assert!(early.failed());
let late = hub.subscribe();
assert!(late.is_closed() && late.failed());
}
#[tokio::test]
async fn a_full_queue_drops_the_oldest_and_reports_the_loss() {
let hub = Hub::new(4);
let mut slow = hub.subscribe();
let mut fast = hub.subscribe();
for i in 0..10 {
hub.publish(&[joined(i)]);
assert_eq!(fast.try_recv(), Some(joined(i)));
}
assert_eq!(slow.recv().await, Some(Event::EventsLost(6)));
for i in 6..10 {
assert_eq!(slow.recv().await, Some(joined(i)));
}
assert_eq!(slow.try_recv(), None);
assert_eq!(fast.try_recv(), None);
}
#[tokio::test]
async fn recv_wakes_when_an_event_arrives() {
let hub = Arc::new(Hub::new(4));
let mut events = hub.subscribe();
let publisher = Arc::clone(&hub);
let task = tokio::spawn(async move { events.recv().await });
tokio::task::yield_now().await;
publisher.publish(&[joined(7)]);
assert_eq!(task.await.unwrap(), Some(joined(7)));
}
#[test]
fn dropped_subscriptions_are_pruned() {
let hub = Hub::new(4);
drop(hub.subscribe());
let _kept = hub.subscribe();
hub.publish(&[joined(0)]);
assert_eq!(lock(&hub.inner).subs.len(), 1);
}
}