use std::sync::Arc;
use dashmap::DashMap;
use parking_lot::Mutex;
use uuid::Uuid;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SubscriptionId(pub u64);
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub enum SubscribeIntent {
Live,
Replay {
from: i64,
},
SnapshotThenLive,
Latest,
Passive,
Ephemeral,
}
#[derive(Debug, Clone)]
pub struct Subscription {
pub id: SubscriptionId,
pub topic: String,
pub intent: SubscribeIntent,
pub cancelled: bool,
}
pub type ConnectionSink = Arc<dyn FanoutSink>;
pub trait FanoutSink: Send + Sync {
fn deliver(&self, frame: bytes::Bytes) -> Result<(), FanoutError>;
fn id(&self) -> u64;
}
#[derive(Debug, thiserror::Error)]
pub enum FanoutError {
#[error("sink closed")]
Closed,
#[error("sink backpressured: queue={queue_bytes}, max={max_bytes}")]
Backpressured {
queue_bytes: usize,
max_bytes: usize,
},
}
pub struct FanoutEngine {
by_topic: DashMap<String, Vec<(SubscriptionId, ConnectionSink)>>,
by_id: DashMap<SubscriptionId, (String, ConnectionSink)>,
seq: Mutex<u64>,
}
impl std::fmt::Debug for FanoutEngine {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FanoutEngine")
.field("subscription_count", &self.by_id.len())
.finish()
}
}
impl Default for FanoutEngine {
fn default() -> Self {
Self::new()
}
}
impl FanoutEngine {
pub fn new() -> Self {
Self {
by_topic: DashMap::new(),
by_id: DashMap::new(),
seq: Mutex::new(0),
}
}
pub fn subscribe(
&self,
topic: &str,
_intent: SubscribeIntent,
sink: ConnectionSink,
) -> SubscriptionId {
let mut seq = self.seq.lock();
*seq += 1;
let id = SubscriptionId(*seq);
drop(seq);
self.by_topic
.entry(topic.to_string())
.or_default()
.push((id, sink.clone()));
self.by_id.insert(id, (topic.to_string(), sink));
id
}
pub fn unsubscribe(&self, id: SubscriptionId) -> Option<String> {
if let Some((_, (topic, _sink))) = self.by_id.remove(&id) {
if let Some(mut list) = self.by_topic.get_mut(&topic) {
list.retain(|(sid, _)| *sid != id);
}
Some(topic)
} else {
None
}
}
pub fn drop_sink(&self, sink_id: u64) -> Vec<String> {
let mut topics = Vec::new();
let ids: Vec<SubscriptionId> = self
.by_id
.iter()
.filter(|kv| kv.value().1.id() == sink_id)
.map(|kv| *kv.key())
.collect();
for id in ids {
if let Some(topic) = self.unsubscribe(id) {
topics.push(topic);
}
}
topics
}
pub fn deliver(&self, topic: &str, frame: bytes::Bytes) -> usize {
let mut ok = 0;
if let Some(list) = self.by_topic.get(topic) {
for (_id, sink) in list.iter() {
if sink.deliver(frame.clone()).is_ok() {
ok += 1;
}
}
}
ok
}
pub fn subscription_count(&self) -> usize {
self.by_id.len()
}
pub fn topic_subscriber_count(&self, topic: &str) -> usize {
self.by_topic.get(topic).map(|l| l.len()).unwrap_or(0)
}
}
pub fn new_sink_id() -> u64 {
let u = Uuid::new_v4();
let bytes = u.as_bytes();
let mut buf = [0u8; 8];
buf.copy_from_slice(&bytes[..8]);
u64::from_le_bytes(buf)
}
pub mod test_sink {
use std::sync::atomic::{AtomicU64, Ordering};
use parking_lot::Mutex;
use super::{FanoutError, FanoutSink};
pub struct CountingSink {
id: u64,
delivered: AtomicU64,
log: Mutex<Vec<Vec<u8>>>,
}
impl CountingSink {
pub fn new(id: u64) -> Self {
Self {
id,
delivered: AtomicU64::new(0),
log: Mutex::new(Vec::new()),
}
}
pub fn count(&self) -> u64 {
self.delivered.load(Ordering::SeqCst)
}
pub fn messages(&self) -> Vec<Vec<u8>> {
self.log.lock().clone()
}
}
impl FanoutSink for CountingSink {
fn deliver(&self, frame: bytes::Bytes) -> Result<(), FanoutError> {
self.delivered.fetch_add(1, Ordering::SeqCst);
self.log.lock().push(frame.to_vec());
Ok(())
}
fn id(&self) -> u64 {
self.id
}
}
}
#[cfg(test)]
mod tests {
use super::test_sink::CountingSink;
use super::*;
#[test]
fn subscribe_and_fanout() {
let fan = FanoutEngine::new();
let s1 = Arc::new(CountingSink::new(1));
let s2 = Arc::new(CountingSink::new(2));
let s3 = Arc::new(CountingSink::new(3));
fan.subscribe("t", SubscribeIntent::Live, s1.clone());
fan.subscribe("t", SubscribeIntent::Live, s2.clone());
fan.subscribe("other", SubscribeIntent::Live, s3.clone());
let frame = bytes::Bytes::from_static(b"hi");
let n = fan.deliver("t", frame);
assert_eq!(n, 2);
assert_eq!(s1.count(), 1);
assert_eq!(s2.count(), 1);
assert_eq!(s3.count(), 0);
}
#[test]
fn unsubscribe_returns_topic() {
let fan = FanoutEngine::new();
let s = Arc::new(CountingSink::new(1));
let id = fan.subscribe("t", SubscribeIntent::Live, s.clone());
let topic = fan.unsubscribe(id);
assert_eq!(topic, Some("t".to_string()));
assert_eq!(fan.deliver("t", bytes::Bytes::from_static(b"x")), 0);
}
#[test]
fn drop_sink_returns_topics() {
let fan = FanoutEngine::new();
let s1 = Arc::new(CountingSink::new(7));
let s2 = Arc::new(CountingSink::new(7));
fan.subscribe("t", SubscribeIntent::Live, s1.clone());
fan.subscribe("u", SubscribeIntent::Live, s2.clone());
let topics = fan.drop_sink(7);
assert_eq!(topics.len(), 2);
assert!(topics.contains(&"t".to_string()));
assert!(topics.contains(&"u".to_string()));
assert_eq!(fan.subscription_count(), 0);
}
#[test]
fn topic_subscriber_count() {
let fan = FanoutEngine::new();
let s = Arc::new(CountingSink::new(1));
fan.subscribe("t", SubscribeIntent::Live, s.clone());
fan.subscribe("t", SubscribeIntent::Live, s.clone());
assert_eq!(fan.topic_subscriber_count("t"), 2);
}
#[test]
fn deliver_records_payload() {
let fan = FanoutEngine::new();
let s = Arc::new(CountingSink::new(1));
fan.subscribe("t", SubscribeIntent::Live, s.clone());
fan.deliver("t", bytes::Bytes::from_static(b"abc"));
assert_eq!(s.messages(), vec![b"abc".to_vec()]);
}
#[test]
fn sink_id_is_unique() {
let a = new_sink_id();
let b = new_sink_id();
assert_ne!(a, b);
}
}