use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{broadcast, mpsc};
use tracing::debug;
use uuid::Uuid;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ChangeType {
Created,
Updated,
Deleted,
Cleared,
}
impl std::fmt::Display for ChangeType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ChangeType::Created => write!(f, "CREATED"),
ChangeType::Updated => write!(f, "UPDATED"),
ChangeType::Deleted => write!(f, "DELETED"),
ChangeType::Cleared => write!(f, "CLEARED"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChangeNotification {
pub key: String,
pub change_type: ChangeType,
pub old_value: Option<String>,
pub new_value: Option<String>,
pub version: u64,
pub timestamp: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SubscriptionId(Uuid);
impl SubscriptionId {
pub fn new() -> Self {
Self(Uuid::new_v4())
}
}
impl Default for SubscriptionId {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub enum NotificationFilter {
All,
Key(String),
KeyPrefix(String),
ChangeType(ChangeType),
And(Vec<NotificationFilter>),
Or(Vec<NotificationFilter>),
}
impl NotificationFilter {
pub fn matches(&self, notification: &ChangeNotification) -> bool {
match self {
NotificationFilter::All => true,
NotificationFilter::Key(key) => notification.key == *key,
NotificationFilter::KeyPrefix(prefix) => notification.key.starts_with(prefix),
NotificationFilter::ChangeType(change_type) => notification.change_type == *change_type,
NotificationFilter::And(filters) => filters.iter().all(|f| f.matches(notification)),
NotificationFilter::Or(filters) => filters.iter().any(|f| f.matches(notification)),
}
}
}
pub struct Subscription {
pub id: SubscriptionId,
pub filter: NotificationFilter,
pub receiver: mpsc::UnboundedReceiver<ChangeNotification>,
}
#[derive(Debug, Clone, Default)]
pub struct NotificationStats {
pub total_notifications_sent: u64,
pub total_subscribers: usize,
pub dropped_notifications: u64,
}
pub struct NotificationBus {
broadcast_tx: broadcast::Sender<ChangeNotification>,
#[allow(clippy::type_complexity)]
subscribers: Arc<
RwLock<
HashMap<
SubscriptionId,
(
NotificationFilter,
mpsc::UnboundedSender<ChangeNotification>,
),
>,
>,
>,
stats: Arc<RwLock<NotificationStats>>,
}
impl NotificationBus {
pub fn new() -> Self {
let (broadcast_tx, _) = broadcast::channel(1000);
Self {
broadcast_tx,
subscribers: Arc::new(RwLock::new(HashMap::new())),
stats: Arc::new(RwLock::new(NotificationStats::default())),
}
}
pub fn subscribe(&self, filter: NotificationFilter) -> Subscription {
let id = SubscriptionId::new();
let (tx, rx) = mpsc::unbounded_channel();
{
let mut subscribers = self.subscribers.write();
subscribers.insert(id, (filter.clone(), tx));
}
{
let mut stats = self.stats.write();
stats.total_subscribers += 1;
}
debug!(
"New subscription created: {:?} with filter: {:?}",
id, filter
);
Subscription {
id,
filter,
receiver: rx,
}
}
pub fn subscribe_all(&self) -> Subscription {
self.subscribe(NotificationFilter::All)
}
pub fn subscribe_key(&self, key: &str) -> Subscription {
self.subscribe(NotificationFilter::Key(key.to_string()))
}
pub fn subscribe_prefix(&self, prefix: &str) -> Subscription {
self.subscribe(NotificationFilter::KeyPrefix(prefix.to_string()))
}
pub fn subscribe_change_type(&self, change_type: ChangeType) -> Subscription {
self.subscribe(NotificationFilter::ChangeType(change_type))
}
pub fn unsubscribe(&self, subscription_id: SubscriptionId) {
let mut subscribers = self.subscribers.write();
if subscribers.remove(&subscription_id).is_some() {
let mut stats = self.stats.write();
stats.total_subscribers = stats.total_subscribers.saturating_sub(1);
debug!("Subscription removed: {:?}", subscription_id);
}
}
pub async fn publish(&self, notification: ChangeNotification) {
{
let mut stats = self.stats.write();
stats.total_notifications_sent += 1;
}
if self.broadcast_tx.send(notification.clone()).is_err() {
}
let subscribers = self.subscribers.read();
let mut dropped_count = 0;
for (filter, sender) in subscribers.values() {
if filter.matches(¬ification) && sender.send(notification.clone()).is_err() {
dropped_count += 1;
}
}
if dropped_count > 0 {
let mut stats = self.stats.write();
stats.dropped_notifications += dropped_count;
debug!(
"Dropped {} notifications due to closed channels",
dropped_count
);
}
debug!(
"Published notification: key={}, type={:?}",
notification.key, notification.change_type
);
}
pub fn broadcast_receiver(&self) -> broadcast::Receiver<ChangeNotification> {
self.broadcast_tx.subscribe()
}
pub fn get_stats(&self) -> NotificationStats {
let stats = self.stats.read();
let subscribers = self.subscribers.read();
NotificationStats {
total_notifications_sent: stats.total_notifications_sent,
total_subscribers: subscribers.len(),
dropped_notifications: stats.dropped_notifications,
}
}
pub fn cleanup_closed_subscribers(&self) {
let mut subscribers = self.subscribers.write();
let initial_count = subscribers.len();
subscribers.retain(|_, (_, sender)| !sender.is_closed());
let removed = initial_count - subscribers.len();
if removed > 0 {
debug!("Cleaned up {} closed subscriber channels", removed);
}
}
pub fn subscriber_count(&self) -> usize {
self.subscribers.read().len()
}
}
impl Default for NotificationBus {
fn default() -> Self {
Self::new()
}
}
pub struct NotificationListener {
subscription: Subscription,
name: String,
}
impl NotificationListener {
pub fn new(subscription: Subscription, name: String) -> Self {
Self { subscription, name }
}
pub async fn listen<F, Fut>(&mut self, mut handler: F)
where
F: FnMut(ChangeNotification) -> Fut,
Fut: std::future::Future<Output = ()>,
{
debug!("Starting notification listener: {}", self.name);
while let Some(notification) = self.subscription.receiver.recv().await {
debug!(
"Listener {} received notification: {:?}",
self.name, notification
);
handler(notification).await;
}
debug!("Notification listener {} stopped", self.name);
}
pub fn subscription_id(&self) -> SubscriptionId {
self.subscription.id
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_notification_bus_basic() {
let bus = NotificationBus::new();
let mut subscription = bus.subscribe_all();
let notification = ChangeNotification {
key: "test_key".to_string(),
change_type: ChangeType::Created,
old_value: None,
new_value: Some("test_value".to_string()),
version: 1,
timestamp: 123456789,
};
bus.publish(notification.clone()).await;
let received = subscription.receiver.recv().await.unwrap();
assert_eq!(received.key, notification.key);
assert_eq!(received.change_type, notification.change_type);
assert_eq!(received.new_value, notification.new_value);
}
#[tokio::test]
async fn test_notification_filtering() {
let bus = NotificationBus::new();
let mut key_subscription = bus.subscribe_key("specific_key");
let mut prefix_subscription = bus.subscribe_prefix("prefix_");
let mut type_subscription = bus.subscribe_change_type(ChangeType::Updated);
let key_notification = ChangeNotification {
key: "specific_key".to_string(),
change_type: ChangeType::Created,
old_value: None,
new_value: Some("value".to_string()),
version: 1,
timestamp: 123456789,
};
let prefix_notification = ChangeNotification {
key: "prefix_test".to_string(),
change_type: ChangeType::Created,
old_value: None,
new_value: Some("value".to_string()),
version: 2,
timestamp: 123456790,
};
let type_notification = ChangeNotification {
key: "any_key".to_string(),
change_type: ChangeType::Updated,
old_value: Some("old".to_string()),
new_value: Some("new".to_string()),
version: 3,
timestamp: 123456791,
};
bus.publish(key_notification.clone()).await;
bus.publish(prefix_notification.clone()).await;
bus.publish(type_notification.clone()).await;
let received = key_subscription.receiver.recv().await.unwrap();
assert_eq!(received.key, "specific_key");
let received = prefix_subscription.receiver.recv().await.unwrap();
assert_eq!(received.key, "prefix_test");
let received = type_subscription.receiver.recv().await.unwrap();
assert_eq!(received.change_type, ChangeType::Updated);
}
#[tokio::test]
async fn test_notification_stats() {
let bus = NotificationBus::new();
let _subscription = bus.subscribe_all();
let notification = ChangeNotification {
key: "test".to_string(),
change_type: ChangeType::Created,
old_value: None,
new_value: Some("value".to_string()),
version: 1,
timestamp: 123456789,
};
bus.publish(notification).await;
let stats = bus.get_stats();
assert_eq!(stats.total_notifications_sent, 1);
assert_eq!(stats.total_subscribers, 1);
}
#[test]
fn test_notification_filter_logic() {
let notification = ChangeNotification {
key: "test_key".to_string(),
change_type: ChangeType::Updated,
old_value: Some("old".to_string()),
new_value: Some("new".to_string()),
version: 1,
timestamp: 123456789,
};
assert!(NotificationFilter::All.matches(¬ification));
assert!(NotificationFilter::Key("test_key".to_string()).matches(¬ification));
assert!(!NotificationFilter::Key("other_key".to_string()).matches(¬ification));
assert!(NotificationFilter::KeyPrefix("test_".to_string()).matches(¬ification));
assert!(!NotificationFilter::KeyPrefix("other_".to_string()).matches(¬ification));
assert!(NotificationFilter::ChangeType(ChangeType::Updated).matches(¬ification));
assert!(!NotificationFilter::ChangeType(ChangeType::Created).matches(¬ification));
let and_filter = NotificationFilter::And(vec![
NotificationFilter::KeyPrefix("test_".to_string()),
NotificationFilter::ChangeType(ChangeType::Updated),
]);
assert!(and_filter.matches(¬ification));
let or_filter = NotificationFilter::Or(vec![
NotificationFilter::Key("wrong_key".to_string()),
NotificationFilter::ChangeType(ChangeType::Updated),
]);
assert!(or_filter.matches(¬ification));
}
}