Fixed simpleBroker panics at called Option::unwrap() on a None value if used with more than one type

src/subscription/simple_broker.rs

@@ -1,21 +1,36 @@
 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() {
-            sender.start_send(msg.clone()).ok();
-        }
+        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();
-        let (tx, rx) = mpsc::unbounded();
-        let id = senders.insert(tx);
-        BrokerStream(id, rx)
+        with_senders::<T, _, _>(|senders| {
+            let (tx, rx) = mpsc::unbounded();
+            let id = senders.0.insert(tx);
+            BrokerStream(id, rx)
+        })
     }
 }