Fixed simpleBroker panics at called `Option::unwrap()` on a `None` value if used with more than one type
use futures::channel::mpsc::{self, UnboundedReceiver, UnboundedSender};
use futures::task::{Context, Poll};
use futures::{Stream, StreamExt};
-use once_cell::sync::OnceCell;
+use once_cell::sync::Lazy;
use serde::export::PhantomData;
use slab::Slab;
-use std::any::Any;
+use std::any::{Any, TypeId};
+use std::collections::HashMap;
use std::pin::Pin;
use std::sync::Mutex;
-struct Senders<T>(Mutex<Slab<UnboundedSender<T>>>);
+static SUBSCRIBERS: Lazy<Mutex<HashMap<TypeId, Box<dyn Any + Send>>>> =
+ Lazy::new(|| Default::default());
+
+struct Senders<T>(Slab<UnboundedSender<T>>);
struct BrokerStream<T: Sync + Send + Clone + 'static>(usize, UnboundedReceiver<T>);
+fn with_senders<T, F, R>(f: F) -> R
+where
+ T: Sync + Send + Clone + 'static,
+ F: FnOnce(&mut Senders<T>) -> R,
+{
+ let mut map = SUBSCRIBERS.lock().unwrap();
+ let senders = map
+ .entry(TypeId::of::<Senders<T>>())
+ .or_insert_with(|| Box::new(Senders::<T>(Default::default())));
+ f(senders.downcast_mut::<Senders<T>>().unwrap())
+}
+
impl<T: Sync + Send + Clone + 'static> Drop for BrokerStream<T> {
fn drop(&mut self) {
- let mut senders = SimpleBroker::<T>::senders().0.lock().unwrap();
- senders.remove(self.0);
+ with_senders::<T, _, _>(|senders| senders.0.remove(self.0));
}
}
@@ -31,25 +46,21 @@ impl<T: Sync + Send + Clone + 'static> Stream for BrokerStream<T> {
pub struct SimpleBroker<T>(PhantomData<T>);
impl<T: Sync + Send + Clone + 'static> SimpleBroker<T> {
- fn senders() -> &'static Senders<T> {
- static SUBSCRIBERS: OnceCell<Box<dyn Any + Send + Sync>> = OnceCell::new();
- let instance = SUBSCRIBERS.get_or_init(|| Box::new(Senders::<T>(Mutex::new(Slab::new()))));
- instance.downcast_ref::<Senders<T>>().unwrap()
- }
-
/// Publish a message that all subscription streams can receive.
pub fn publish(msg: T) {
- let mut senders = Self::senders().0.lock().unwrap();
- for (_, sender) in senders.iter_mut() {
+ with_senders::<T, _, _>(|senders| {
+ for (_, sender) in senders.0.iter_mut() {
sender.start_send(msg.clone()).ok();
}
+ });
}
/// Subscribe to the message of the specified type and returns a `Stream`.
pub fn subscribe() -> impl Stream<Item = T> {
- let mut senders = Self::senders().0.lock().unwrap();
+ with_senders::<T, _, _>(|senders| {
let (tx, rx) = mpsc::unbounded();
- let id = senders.insert(tx);
+ let id = senders.0.insert(tx);
BrokerStream(id, rx)
+ })
}
}
