use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use aion_core::{ClusterEvent, ClusterEventMeta};
use futures::stream::{self, BoxStream};
use tokio::sync::broadcast;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ClusterStreamLagged {
pub skipped: u64,
}
#[derive(Clone, Debug)]
pub struct ClusterEventPublisher {
events: broadcast::Sender<ClusterEvent>,
next_seq: Arc<AtomicU64>,
}
impl ClusterEventPublisher {
#[must_use]
pub fn new(capacity: std::num::NonZeroUsize) -> Self {
let (events, _receiver) = broadcast::channel(capacity.get());
Self {
events,
next_seq: Arc::new(AtomicU64::new(1)),
}
}
pub fn emit<F>(&self, build: F) -> ClusterEvent
where
F: FnOnce(ClusterEventMeta) -> ClusterEvent,
{
let meta = ClusterEventMeta {
cluster_seq: self.next_seq.fetch_add(1, Ordering::SeqCst),
observed_at: chrono::Utc::now(),
};
let event = build(meta);
let send_result = self.events.send(event.clone());
drop(send_result);
event
}
#[must_use]
pub fn subscribe(
&self,
after_seq: u64,
) -> BoxStream<'static, Result<ClusterEvent, ClusterStreamLagged>> {
let receiver = self.events.subscribe();
Box::pin(stream::unfold(
(receiver, after_seq),
|(mut receiver, after_seq)| async move {
loop {
match receiver.recv().await {
Ok(event) => {
if event_seq(&event) > after_seq {
return Some((Ok(event), (receiver, after_seq)));
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
return Some((
Err(ClusterStreamLagged { skipped }),
(receiver, after_seq),
));
}
Err(broadcast::error::RecvError::Closed) => return None,
}
}
},
))
}
#[must_use]
pub fn current_seq(&self) -> u64 {
self.next_seq.load(Ordering::SeqCst).saturating_sub(1)
}
}
fn event_seq(event: &ClusterEvent) -> u64 {
cluster_event_meta(event).cluster_seq
}
#[must_use]
pub fn cluster_event_meta(event: &ClusterEvent) -> &ClusterEventMeta {
match event {
ClusterEvent::PeerAdded { meta, .. }
| ClusterEvent::PeerConnected { meta, .. }
| ClusterEvent::PeerDisconnected { meta, .. }
| ClusterEvent::ShardAdopted { meta, .. }
| ClusterEvent::ShardAdoptionFailed { meta, .. }
| ClusterEvent::ShardAdoptionSkipped { meta, .. }
| ClusterEvent::WorkerConnected { meta, .. }
| ClusterEvent::WorkerDisconnected { meta, .. }
| ClusterEvent::SupervisorStarted { meta, .. }
| ClusterEvent::SupervisorStopped { meta, .. }
| ClusterEvent::NamespaceCreated { meta, .. }
| ClusterEvent::NamespacePlacementChanged { meta, .. }
| ClusterEvent::NamespaceQuotaState { meta, .. } => meta,
}
}
#[cfg(test)]
mod tests {
use std::num::NonZeroUsize;
use aion_core::ClusterEvent;
use futures::StreamExt;
use super::*;
fn capacity(value: usize) -> Result<NonZeroUsize, Box<dyn std::error::Error>> {
NonZeroUsize::new(value).ok_or_else(|| "capacity must be non-zero".into())
}
fn supervisor_started(meta: ClusterEventMeta) -> ClusterEvent {
ClusterEvent::SupervisorStarted {
meta,
node: "node-1@127.0.0.1".to_owned(),
}
}
#[tokio::test]
async fn emit_stamps_monotonic_increasing_seq() -> Result<(), Box<dyn std::error::Error>> {
let publisher = ClusterEventPublisher::new(capacity(8)?);
let mut subscription = publisher.subscribe(0);
let first = publisher.emit(supervisor_started);
let second = publisher.emit(supervisor_started);
assert_eq!(cluster_event_meta(&first).cluster_seq, 1);
assert_eq!(cluster_event_meta(&second).cluster_seq, 2);
let received_first = subscription
.next()
.await
.ok_or("missing first")?
.map_err(|lag| format!("unexpected lag: {lag:?}"))?;
let received_second = subscription
.next()
.await
.ok_or("missing second")?
.map_err(|lag| format!("unexpected lag: {lag:?}"))?;
assert_eq!(cluster_event_meta(&received_first).cluster_seq, 1);
assert_eq!(cluster_event_meta(&received_second).cluster_seq, 2);
Ok(())
}
#[tokio::test]
async fn after_seq_suppresses_already_applied_splice_deltas()
-> Result<(), Box<dyn std::error::Error>> {
let publisher = ClusterEventPublisher::new(capacity(8)?);
let mut subscription = publisher.subscribe(2);
for _ in 0..3 {
publisher.emit(supervisor_started);
}
let survivor = subscription
.next()
.await
.ok_or("missing survivor")?
.map_err(|lag| format!("unexpected lag: {lag:?}"))?;
assert_eq!(
cluster_event_meta(&survivor).cluster_seq,
3,
"deltas at or below after_seq must be suppressed at the splice seam"
);
Ok(())
}
#[tokio::test]
async fn lagged_subscriber_yields_typed_skip_count() -> Result<(), Box<dyn std::error::Error>> {
let publisher = ClusterEventPublisher::new(capacity(2)?);
let mut subscription = publisher.subscribe(0);
for _ in 0..5 {
publisher.emit(supervisor_started);
}
let lagged = subscription.next().await.ok_or("missing lag item")?;
assert_eq!(lagged, Err(ClusterStreamLagged { skipped: 3 }));
Ok(())
}
#[tokio::test]
async fn emit_with_no_subscribers_is_not_an_error() -> Result<(), Box<dyn std::error::Error>> {
let publisher = ClusterEventPublisher::new(capacity(2)?);
let event = publisher.emit(supervisor_started);
assert_eq!(cluster_event_meta(&event).cluster_seq, 1);
assert_eq!(publisher.current_seq(), 1);
Ok(())
}
}