|
@@ -1,339 +0,0 @@
|
|
|
-//! Publish–subscribe pattern.
|
|
|
|
|
-//!
|
|
|
|
|
-//! This is a generic implementation for
|
|
|
|
|
-//! [NUT-17](<https://github.com/cashubtc/nuts/blob/main/17.md>) with a type
|
|
|
|
|
-//! agnostic Publish-subscribe manager.
|
|
|
|
|
-//!
|
|
|
|
|
-//! The manager has a method for subscribers to subscribe to events with a
|
|
|
|
|
-//! generic type that must be converted to a vector of indexes.
|
|
|
|
|
-//!
|
|
|
|
|
-//! Events are also generic that should implement the `Indexable` trait.
|
|
|
|
|
-use std::cmp::Ordering;
|
|
|
|
|
-use std::collections::{BTreeMap, HashSet};
|
|
|
|
|
-use std::fmt::Debug;
|
|
|
|
|
-use std::ops::{Deref, DerefMut};
|
|
|
|
|
-use std::sync::atomic::{self, AtomicUsize};
|
|
|
|
|
-use std::sync::Arc;
|
|
|
|
|
-
|
|
|
|
|
-pub use cdk_common::pub_sub::index::{Index, Indexable, SubscriptionGlobalId};
|
|
|
|
|
-use cdk_common::pub_sub::OnNewSubscription;
|
|
|
|
|
-pub use cdk_common::pub_sub::SubId;
|
|
|
|
|
-use tokio::sync::{mpsc, RwLock};
|
|
|
|
|
-use tokio::task::JoinHandle;
|
|
|
|
|
-
|
|
|
|
|
-type IndexTree<T, I> = Arc<RwLock<BTreeMap<Index<I>, mpsc::Sender<(SubId, T)>>>>;
|
|
|
|
|
-
|
|
|
|
|
-/// Default size of the remove channel
|
|
|
|
|
-pub const DEFAULT_REMOVE_SIZE: usize = 10_000;
|
|
|
|
|
-
|
|
|
|
|
-/// Default channel size for subscription buffering
|
|
|
|
|
-pub const DEFAULT_CHANNEL_SIZE: usize = 10;
|
|
|
|
|
-
|
|
|
|
|
-/// Subscription manager
|
|
|
|
|
-///
|
|
|
|
|
-/// This object keep track of all subscription listener and it is also
|
|
|
|
|
-/// responsible for broadcasting events to all listeners
|
|
|
|
|
-///
|
|
|
|
|
-/// The content of the notification is not relevant to this scope and it is up
|
|
|
|
|
-/// to the application, therefore the generic T is used instead of a specific
|
|
|
|
|
-/// type
|
|
|
|
|
-pub struct Manager<T, I, F>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable<Type = I> + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
|
|
|
|
|
- F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- indexes: IndexTree<T, I>,
|
|
|
|
|
- on_new_subscription: Option<Arc<F>>,
|
|
|
|
|
- unsubscription_sender: mpsc::Sender<(SubId, Vec<Index<I>>)>,
|
|
|
|
|
- active_subscriptions: Arc<AtomicUsize>,
|
|
|
|
|
- background_subscription_remover: Option<JoinHandle<()>>,
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-impl<T, I, F> Default for Manager<T, I, F>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable<Type = I> + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
|
|
|
|
|
- F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- fn default() -> Self {
|
|
|
|
|
- let (sender, receiver) = mpsc::channel(DEFAULT_REMOVE_SIZE);
|
|
|
|
|
- let active_subscriptions: Arc<AtomicUsize> = Default::default();
|
|
|
|
|
- let storage: IndexTree<T, I> = Arc::new(Default::default());
|
|
|
|
|
-
|
|
|
|
|
- Self {
|
|
|
|
|
- background_subscription_remover: Some(tokio::spawn(Self::remove_subscription(
|
|
|
|
|
- receiver,
|
|
|
|
|
- storage.clone(),
|
|
|
|
|
- active_subscriptions.clone(),
|
|
|
|
|
- ))),
|
|
|
|
|
- on_new_subscription: None,
|
|
|
|
|
- unsubscription_sender: sender,
|
|
|
|
|
- active_subscriptions,
|
|
|
|
|
- indexes: storage,
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-impl<T, I, F> From<F> for Manager<T, I, F>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable<Type = I> + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
|
|
|
|
|
- F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- fn from(value: F) -> Self {
|
|
|
|
|
- let mut manager: Self = Default::default();
|
|
|
|
|
- manager.on_new_subscription = Some(Arc::new(value));
|
|
|
|
|
- manager
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-impl<T, I, F> Manager<T, I, F>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable<Type = I> + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: PartialOrd + Clone + Debug + Ord + Send + Sync + 'static,
|
|
|
|
|
- F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- #[inline]
|
|
|
|
|
- /// Broadcast an event to all listeners
|
|
|
|
|
- ///
|
|
|
|
|
- /// This function takes an Arc to the storage struct, the event_id, the kind
|
|
|
|
|
- /// and the vent to broadcast
|
|
|
|
|
- async fn broadcast_impl(storage: &IndexTree<T, I>, event: T) {
|
|
|
|
|
- let index_storage = storage.read().await;
|
|
|
|
|
- let mut sent = HashSet::new();
|
|
|
|
|
- for index in event.to_indexes() {
|
|
|
|
|
- for (key, sender) in index_storage.range(index.clone()..) {
|
|
|
|
|
- if index.cmp_prefix(key) != Ordering::Equal {
|
|
|
|
|
- break;
|
|
|
|
|
- }
|
|
|
|
|
- let sub_id = key.unique_id();
|
|
|
|
|
- if sent.contains(&sub_id) {
|
|
|
|
|
- continue;
|
|
|
|
|
- }
|
|
|
|
|
- sent.insert(sub_id);
|
|
|
|
|
- let _ = sender.try_send((key.into(), event.clone()));
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Broadcasts an event to all listeners
|
|
|
|
|
- ///
|
|
|
|
|
- /// This public method will not block the caller, it will spawn a new task
|
|
|
|
|
- /// instead
|
|
|
|
|
- pub fn broadcast(&self, event: T) {
|
|
|
|
|
- let storage = self.indexes.clone();
|
|
|
|
|
- tokio::spawn(async move {
|
|
|
|
|
- Self::broadcast_impl(&storage, event).await;
|
|
|
|
|
- });
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Broadcasts an event to all listeners
|
|
|
|
|
- ///
|
|
|
|
|
- /// This method is async and will await for the broadcast to be completed
|
|
|
|
|
- pub async fn broadcast_async(&self, event: T) {
|
|
|
|
|
- Self::broadcast_impl(&self.indexes, event).await;
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Specific of the subscription, this is the abstraction between `subscribe` and `try_subscribe`
|
|
|
|
|
- #[inline(always)]
|
|
|
|
|
- async fn subscribe_inner(
|
|
|
|
|
- &self,
|
|
|
|
|
- sub_id: SubId,
|
|
|
|
|
- indexes: Vec<Index<I>>,
|
|
|
|
|
- ) -> ActiveSubscription<T, I> {
|
|
|
|
|
- let (sender, receiver) = mpsc::channel(10);
|
|
|
|
|
-
|
|
|
|
|
- let mut index_storage = self.indexes.write().await;
|
|
|
|
|
- // Subscribe to events as soon as possible
|
|
|
|
|
- for index in indexes.clone() {
|
|
|
|
|
- index_storage.insert(index, sender.clone());
|
|
|
|
|
- }
|
|
|
|
|
- drop(index_storage);
|
|
|
|
|
-
|
|
|
|
|
- if let Some(on_new_subscription) = self.on_new_subscription.clone() {
|
|
|
|
|
- // After we're subscribed already, fetch the current status of matching events. It is
|
|
|
|
|
- // down in another thread to return right away
|
|
|
|
|
- let indexes_for_worker = indexes.clone();
|
|
|
|
|
- let sub_id_for_worker = sub_id.clone();
|
|
|
|
|
- tokio::spawn(async move {
|
|
|
|
|
- match on_new_subscription
|
|
|
|
|
- .on_new_subscription(
|
|
|
|
|
- &indexes_for_worker
|
|
|
|
|
- .iter()
|
|
|
|
|
- .map(|x| x.deref())
|
|
|
|
|
- .collect::<Vec<_>>(),
|
|
|
|
|
- )
|
|
|
|
|
- .await
|
|
|
|
|
- {
|
|
|
|
|
- Ok(events) => {
|
|
|
|
|
- for event in events {
|
|
|
|
|
- let _ = sender.try_send((sub_id_for_worker.clone(), event));
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- Err(err) => {
|
|
|
|
|
- tracing::info!(
|
|
|
|
|
- "Failed to get initial state for subscription: {:?}, {}",
|
|
|
|
|
- sub_id_for_worker,
|
|
|
|
|
- err
|
|
|
|
|
- );
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
- });
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- self.active_subscriptions
|
|
|
|
|
- .fetch_add(1, atomic::Ordering::Relaxed);
|
|
|
|
|
-
|
|
|
|
|
- ActiveSubscription {
|
|
|
|
|
- sub_id,
|
|
|
|
|
- receiver,
|
|
|
|
|
- indexes,
|
|
|
|
|
- drop: self.unsubscription_sender.clone(),
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Try to subscribe to a specific event
|
|
|
|
|
- pub async fn try_subscribe<P>(&self, params: P) -> Result<ActiveSubscription<T, I>, P::Error>
|
|
|
|
|
- where
|
|
|
|
|
- P: AsRef<SubId> + TryInto<Vec<Index<I>>>,
|
|
|
|
|
- {
|
|
|
|
|
- Ok(self
|
|
|
|
|
- .subscribe_inner(params.as_ref().clone(), params.try_into()?)
|
|
|
|
|
- .await)
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Subscribe to a specific event
|
|
|
|
|
- pub async fn subscribe<P>(&self, params: P) -> ActiveSubscription<T, I>
|
|
|
|
|
- where
|
|
|
|
|
- P: AsRef<SubId> + Into<Vec<Index<I>>>,
|
|
|
|
|
- {
|
|
|
|
|
- self.subscribe_inner(params.as_ref().clone(), params.into())
|
|
|
|
|
- .await
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Return number of active subscriptions
|
|
|
|
|
- pub fn active_subscriptions(&self) -> usize {
|
|
|
|
|
- self.active_subscriptions.load(atomic::Ordering::SeqCst)
|
|
|
|
|
- }
|
|
|
|
|
-
|
|
|
|
|
- /// Task to remove dropped subscriptions from the storage struct
|
|
|
|
|
- ///
|
|
|
|
|
- /// This task will run in the background (and will be dropped when the [`Manager`]
|
|
|
|
|
- /// is) and will remove subscriptions from the storage struct it is dropped.
|
|
|
|
|
- async fn remove_subscription(
|
|
|
|
|
- mut receiver: mpsc::Receiver<(SubId, Vec<Index<I>>)>,
|
|
|
|
|
- storage: IndexTree<T, I>,
|
|
|
|
|
- active_subscriptions: Arc<AtomicUsize>,
|
|
|
|
|
- ) {
|
|
|
|
|
- while let Some((sub_id, indexes)) = receiver.recv().await {
|
|
|
|
|
- tracing::info!("Removing subscription: {}", *sub_id);
|
|
|
|
|
-
|
|
|
|
|
- active_subscriptions.fetch_sub(1, atomic::Ordering::AcqRel);
|
|
|
|
|
-
|
|
|
|
|
- let mut index_storage = storage.write().await;
|
|
|
|
|
- for key in indexes {
|
|
|
|
|
- index_storage.remove(&key);
|
|
|
|
|
- }
|
|
|
|
|
- drop(index_storage);
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-/// Manager goes out of scope, stop all background tasks
|
|
|
|
|
-impl<T, I, F> Drop for Manager<T, I, F>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable<Type = I> + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
|
|
|
|
|
- F: OnNewSubscription<Index = I, Event = T> + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- fn drop(&mut self) {
|
|
|
|
|
- if let Some(handler) = self.background_subscription_remover.take() {
|
|
|
|
|
- handler.abort();
|
|
|
|
|
- }
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-/// Active Subscription
|
|
|
|
|
-///
|
|
|
|
|
-/// This struct is a wrapper around the `mpsc::Receiver<Event>` and it also used
|
|
|
|
|
-/// to keep track of the subscription itself. When this struct goes out of
|
|
|
|
|
-/// scope, it will notify the Manager about it, so it can be removed from the
|
|
|
|
|
-/// list of active listeners
|
|
|
|
|
-pub struct ActiveSubscription<T, I>
|
|
|
|
|
-where
|
|
|
|
|
- T: Send + Sync,
|
|
|
|
|
- I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- /// The subscription ID
|
|
|
|
|
- pub sub_id: SubId,
|
|
|
|
|
- indexes: Vec<Index<I>>,
|
|
|
|
|
- receiver: mpsc::Receiver<(SubId, T)>,
|
|
|
|
|
- drop: mpsc::Sender<(SubId, Vec<Index<I>>)>,
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-impl<T, I> Deref for ActiveSubscription<T, I>
|
|
|
|
|
-where
|
|
|
|
|
- T: Send + Sync,
|
|
|
|
|
- I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- type Target = mpsc::Receiver<(SubId, T)>;
|
|
|
|
|
-
|
|
|
|
|
- fn deref(&self) -> &Self::Target {
|
|
|
|
|
- &self.receiver
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-impl<T, I> DerefMut for ActiveSubscription<T, I>
|
|
|
|
|
-where
|
|
|
|
|
- T: Indexable + Clone + Send + Sync + 'static,
|
|
|
|
|
- I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- fn deref_mut(&mut self) -> &mut Self::Target {
|
|
|
|
|
- &mut self.receiver
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-/// The ActiveSubscription is Drop out of scope, notify the Manager about it, so
|
|
|
|
|
-/// it can be removed from the list of active listeners
|
|
|
|
|
-///
|
|
|
|
|
-/// Having this in place, we can avoid memory leaks and also makes it super
|
|
|
|
|
-/// simple to implement the Unsubscribe method
|
|
|
|
|
-impl<T, I> Drop for ActiveSubscription<T, I>
|
|
|
|
|
-where
|
|
|
|
|
- T: Send + Sync,
|
|
|
|
|
- I: Clone + Debug + PartialOrd + Ord + Send + Sync + 'static,
|
|
|
|
|
-{
|
|
|
|
|
- fn drop(&mut self) {
|
|
|
|
|
- let _ = self
|
|
|
|
|
- .drop
|
|
|
|
|
- .try_send((self.sub_id.clone(), self.indexes.drain(..).collect()));
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|
|
|
-
|
|
|
|
|
-#[cfg(test)]
|
|
|
|
|
-mod test {
|
|
|
|
|
- use tokio::sync::mpsc;
|
|
|
|
|
-
|
|
|
|
|
- use super::*;
|
|
|
|
|
-
|
|
|
|
|
- #[test]
|
|
|
|
|
- fn test_active_subscription_drop() {
|
|
|
|
|
- let (tx, rx) = mpsc::channel::<(SubId, ())>(10);
|
|
|
|
|
- let sub_id = SubId::from("test_sub_id");
|
|
|
|
|
- let indexes: Vec<Index<String>> = vec![Index::from(("test".to_string(), sub_id.clone()))];
|
|
|
|
|
- let (drop_tx, mut drop_rx) = mpsc::channel(10);
|
|
|
|
|
-
|
|
|
|
|
- {
|
|
|
|
|
- let _active_subscription = ActiveSubscription {
|
|
|
|
|
- sub_id: sub_id.clone(),
|
|
|
|
|
- indexes,
|
|
|
|
|
- receiver: rx,
|
|
|
|
|
- drop: drop_tx,
|
|
|
|
|
- };
|
|
|
|
|
- // When it goes out of scope, it should notify
|
|
|
|
|
- }
|
|
|
|
|
- assert_eq!(drop_rx.try_recv().unwrap().0, sub_id); // it should have notified
|
|
|
|
|
- assert!(tx.try_send(("foo".into(), ())).is_err()); // subscriber is dropped
|
|
|
|
|
- }
|
|
|
|
|
-}
|
|
|