subscriptions: Share watcher among subscriptions of same filter

Author: leoyvens

core/src/subgraph/registrar.rs

@@ -136,7 +136,7 @@ where
         let logger = self.logger.clone();
 
         self.subscription_manager
-            .subscribe(vec![SubscriptionFilter::Assignment])
+            .subscribe(FromIterator::from_iter([SubscriptionFilter::Assignment]))
             .map_err(|()| anyhow!("Entity change stream failed"))
             .map(|event| {
                 // We're only interested in the SubgraphDeploymentAssignment change; we

graph/src/components/store.rs

@@ -588,7 +588,7 @@ impl StoreEvent {
         self
     }
 
-    pub fn matches(&self, filters: &Vec<SubscriptionFilter>) -> bool {
+    pub fn matches(&self, filters: &BTreeSet<SubscriptionFilter>) -> bool {
         self.changes
             .iter()
             .any(|change| filters.iter().any(|filter| filter.matches(change)))
@@ -649,7 +649,7 @@ where
     /// Filter a `StoreEventStream` by subgraph and entity. Only events that have
     /// at least one change to one of the given (subgraph, entity) combinations
     /// will be delivered by the filtered stream.
-    pub fn filter_by_entities(self, filters: Vec<SubscriptionFilter>) -> StoreEventStreamBox {
+    pub fn filter_by_entities(self, filters: BTreeSet<SubscriptionFilter>) -> StoreEventStreamBox {
         let source = self.source.filter(move |event| event.matches(&filters));
 
         StoreEventStream::new(Box::new(source))
@@ -852,10 +852,10 @@ pub trait SubscriptionManager: Send + Sync + 'static {
     /// Subscribe to changes for specific subgraphs and entities.
     ///
     /// Returns a stream of store events that match the input arguments.
-    fn subscribe(&self, entities: Vec<SubscriptionFilter>) -> StoreEventStreamBox;
+    fn subscribe(&self, entities: BTreeSet<SubscriptionFilter>) -> StoreEventStreamBox;
 
     /// If the payload is not required, use for a more efficient subscription mechanism backed by a watcher.
-    fn subscribe_no_payload(&self, entities: Vec<SubscriptionFilter>) -> UnitStream;
+    fn subscribe_no_payload(&self, entities: BTreeSet<SubscriptionFilter>) -> UnitStream;
 }
 
 /// An internal identifer for the specific instance of a deployment. The

graph/src/data/store/mod.rs

@@ -26,7 +26,7 @@ pub mod scalar;
 pub mod ethereum;
 
 /// Filter subscriptions
-#[derive(Clone)]
+#[derive(Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub enum SubscriptionFilter {
     /// Receive updates about all entities from the given deployment of the
     /// given type

graphql/src/store/query.rs

@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque};
 use std::mem::discriminant;
 
 use graph::prelude::*;
@@ -295,7 +295,7 @@ pub fn collect_entities_from_query_field(
     schema: &s::Document,
     object_type: &s::ObjectType,
     field: &q::Field,
-) -> Vec<SubscriptionFilter> {
+) -> BTreeSet<SubscriptionFilter> {
     // Output entities
     let mut entities = HashSet::new();
 

node/src/manager/commands/listen.rs

@@ -1,5 +1,6 @@
-use std::io::Write;
+use std::iter::FromIterator;
 use std::sync::Arc;
+use std::{collections::BTreeSet, io::Write};
 
 use futures::compat::Future01CompatExt;
 //use futures::future;
@@ -13,7 +14,7 @@ use crate::manager::deployment;
 
 async fn listen(
     mgr: Arc<SubscriptionManager>,
-    filter: Vec<SubscriptionFilter>,
+    filter: BTreeSet<SubscriptionFilter>,
 ) -> Result<(), Error> {
     let events = mgr.subscribe(filter);
     println!("press ctrl-c to stop");
@@ -41,7 +42,11 @@ async fn listen(
 
 pub async fn assignments(mgr: Arc<SubscriptionManager>) -> Result<(), Error> {
     println!("waiting for assignment events");
-    listen(mgr, vec![SubscriptionFilter::Assignment]).await?;
+    listen(
+        mgr,
+        FromIterator::from_iter([SubscriptionFilter::Assignment]),
+    )
+    .await?;
 
     Ok(())
 }
@@ -52,7 +57,7 @@ pub async fn entities(
     entity_types: Vec<String>,
 ) -> Result<(), Error> {
     let deployment = deployment::as_hash(deployment)?;
-    let filter: Vec<_> = entity_types
+    let filter = entity_types
         .into_iter()
         .map(|et| SubscriptionFilter::Entities(deployment.clone(), EntityType::new(et)))
         .collect();

node/src/manager/mod.rs

@@ -1,3 +1,5 @@
+use std::collections::BTreeSet;
+
 use graph::{
     components::store::{SubscriptionManager, UnitStream},
     prelude::{StoreEventStreamBox, SubscriptionFilter},
@@ -12,11 +14,11 @@ mod display;
 pub struct PanicSubscriptionManager;
 
 impl SubscriptionManager for PanicSubscriptionManager {
-    fn subscribe(&self, _: Vec<SubscriptionFilter>) -> StoreEventStreamBox {
+    fn subscribe(&self, _: BTreeSet<SubscriptionFilter>) -> StoreEventStreamBox {
         panic!("we were never meant to call `subscribe`");
     }
 
-    fn subscribe_no_payload(&self, _: Vec<SubscriptionFilter>) -> UnitStream {
+    fn subscribe_no_payload(&self, _: BTreeSet<SubscriptionFilter>) -> UnitStream {
         panic!("we were never meant to call `subscribe_no_payload`");
     }
 }

store/postgres/src/store_events.rs

@@ -1,5 +1,7 @@
 use futures03::TryStreamExt;
+use graph::parking_lot::Mutex;
 use graph::tokio_stream::wrappers::ReceiverStream;
+use std::collections::BTreeSet;
 use std::sync::{atomic::Ordering, Arc, RwLock};
 use std::{collections::HashMap, sync::atomic::AtomicUsize};
 use tokio::sync::mpsc::{channel, Sender};
@@ -84,38 +86,46 @@ impl StoreEventListener {
     }
 }
 
-#[async_trait]
-trait EventSink: Send + Sync {
-    async fn send(&self, event: Arc<StoreEvent>) -> Result<(), Error>;
-    fn is_closed(&self) -> bool;
+struct Watcher<T> {
+    sender: Arc<watch::Sender<T>>,
+    receiver: watch::Receiver<T>,
 }
 
-#[async_trait]
-impl EventSink for Sender<Arc<StoreEvent>> {
-    async fn send(&self, event: Arc<StoreEvent>) -> Result<(), Error> {
-        Ok(self.send(event).await?)
+impl<T: Clone + Debug + Send + Sync + 'static> Watcher<T> {
+    fn new(init: T) -> Self {
+        let (sender, receiver) = watch::channel(init);
+        Watcher {
+            sender: Arc::new(sender),
+            receiver,
+        }
     }
 
-    fn is_closed(&self) -> bool {
-        self.is_closed()
+    fn send(&self, v: T) {
+        // Unwrap: `self` holds a receiver.
+        self.sender.send(v).unwrap()
     }
-}
 
-#[async_trait]
-impl EventSink for watch::Sender<()> {
-    async fn send(&self, _event: Arc<StoreEvent>) -> Result<(), Error> {
-        Ok(self.send(())?)
+    fn stream(&self) -> Box<dyn futures03::Stream<Item = T> + Unpin + Send + Sync> {
+        Box::new(tokio_stream::wrappers::WatchStream::new(
+            self.receiver.clone(),
+        ))
     }
 
-    fn is_closed(&self) -> bool {
-        self.is_closed()
+    /// Outstanding receivers returned from `Self::stream`.
+    fn receiver_count(&self) -> usize {
+        // Do not count the internal receiver.
+        self.sender.receiver_count() - 1
     }
 }
 
 /// Manage subscriptions to the `StoreEvent` stream. Keep a list of
 /// currently active subscribers and forward new events to each of them
 pub struct SubscriptionManager {
-    subscriptions: Arc<RwLock<HashMap<String, (Arc<Vec<SubscriptionFilter>>, Arc<dyn EventSink>)>>>,
+    // These are more efficient since only one entry is stored per filter.
+    subscriptions_no_payload: Arc<Mutex<HashMap<BTreeSet<SubscriptionFilter>, Watcher<()>>>>,
+
+    subscriptions:
+        Arc<RwLock<HashMap<String, (Arc<BTreeSet<SubscriptionFilter>>, Sender<Arc<StoreEvent>>)>>>,
 
     /// Keep the notification listener alive
     listener: StoreEventListener,
@@ -126,6 +136,7 @@ impl SubscriptionManager {
         let (listener, store_events) = StoreEventListener::new(logger, postgres_url, registry);
 
         let mut manager = SubscriptionManager {
+            subscriptions_no_payload: Arc::new(Mutex::new(HashMap::new())),
             subscriptions: Arc::new(RwLock::new(HashMap::new())),
             listener,
         };
@@ -146,61 +157,101 @@ impl SubscriptionManager {
         &self,
         store_events: Box<dyn Stream<Item = StoreEvent, Error = ()> + Send>,
     ) {
-        let subscriptions = self.subscriptions.clone();
+        let subscriptions = self.subscriptions.cheap_clone();
+        let subscriptions_no_payload = self.subscriptions_no_payload.cheap_clone();
         let mut store_events = store_events.compat();
 
         // This channel is constantly receiving things and there are locks involved,
         // so it's best to use a blocking task.
         graph::spawn_blocking(async move {
             while let Some(Ok(event)) = store_events.next().await {
-                let senders = subscriptions.read().unwrap().clone();
                 let event = Arc::new(event);
 
-                // Write change to all matching subscription streams; remove subscriptions
-                // whose receiving end has been dropped
-                for (id, (_, sender)) in senders
-                    .iter()
-                    .filter(|(_, (filter, _))| event.matches(filter))
+                // Send to `subscriptions`.
                 {
-                    if sender.send(event.cheap_clone()).await.is_err() {
-                        // Receiver was dropped
-                        subscriptions.write().unwrap().remove(id);
+                    let senders = subscriptions.read().unwrap().clone();
+
+                    // Write change to all matching subscription streams; remove subscriptions
+                    // whose receiving end has been dropped
+                    for (id, (_, sender)) in senders
+                        .iter()
+                        .filter(|(_, (filter, _))| event.matches(filter))
+                    {
+                        if sender.send(event.cheap_clone()).await.is_err() {
+                            // Receiver was dropped
+                            subscriptions.write().unwrap().remove(id);
+                        }
+                    }
+                }
+
+                // Send to `subscriptions_no_payload`.
+                {
+                    let watchers = subscriptions_no_payload.lock();
+
+                    // Write change to all matching subscription streams
+                    for (_, watcher) in watchers.iter().filter(|(filter, _)| event.matches(filter))
+                    {
+                        watcher.send(());
                     }
                 }
             }
         });
     }
 
     fn periodically_clean_up_stale_subscriptions(&self) {
-        let subscriptions = self.subscriptions.clone();
+        let subscriptions = self.subscriptions.cheap_clone();
+        let subscriptions_no_payload = self.subscriptions_no_payload.cheap_clone();
 
         // Clean up stale subscriptions every 5s
         graph::spawn(async move {
             let mut interval = tokio::time::interval(Duration::from_secs(5));
             loop {
                 interval.tick().await;
-                let mut subscriptions = subscriptions.write().unwrap();
-
-                // Obtain IDs of subscriptions whose receiving end has gone
-                let stale_ids = subscriptions
-                    .iter_mut()
-                    .filter_map(|(id, (_, sender))| match sender.is_closed() {
-                        true => Some(id.clone()),
-                        false => None,
-                    })
-                    .collect::<Vec<_>>();
-
-                // Remove all stale subscriptions
-                for id in stale_ids {
-                    subscriptions.remove(&id);
+
+                // Cleanup `subscriptions`.
+                {
+                    let mut subscriptions = subscriptions.write().unwrap();
+
+                    // Obtain IDs of subscriptions whose receiving end has gone
+                    let stale_ids = subscriptions
+                        .iter_mut()
+                        .filter_map(|(id, (_, sender))| match sender.is_closed() {
+                            true => Some(id.clone()),
+                            false => None,
+                        })
+                        .collect::<Vec<_>>();
+
+                    // Remove all stale subscriptions
+                    for id in stale_ids {
+                        subscriptions.remove(&id);
+                    }
+                }
+
+                // Cleanup `subscriptions_no_payload`.
+                {
+                    let mut subscriptions = subscriptions_no_payload.lock();
+
+                    // Obtain IDs of subscriptions whose receiving end has gone
+                    let stale_ids = subscriptions
+                        .iter_mut()
+                        .filter_map(|(id, watcher)| match watcher.receiver_count() == 0 {
+                            true => Some(id.clone()),
+                            false => None,
+                        })
+                        .collect::<Vec<_>>();
+
+                    // Remove all stale subscriptions
+                    for id in stale_ids {
+                        subscriptions.remove(&id);
+                    }
                 }
             }
         });
     }
 }
 
 impl SubscriptionManagerTrait for SubscriptionManager {
-    fn subscribe(&self, entities: Vec<SubscriptionFilter>) -> StoreEventStreamBox {
+    fn subscribe(&self, entities: BTreeSet<SubscriptionFilter>) -> StoreEventStreamBox {
         let id = Uuid::new_v4().to_string();
 
         // Prepare the new subscription by creating a channel and a subscription object
@@ -210,23 +261,18 @@ impl SubscriptionManagerTrait for SubscriptionManager {
         self.subscriptions
             .write()
             .unwrap()
-            .insert(id, (Arc::new(entities.clone()), Arc::new(sender)));
+            .insert(id, (Arc::new(entities.clone()), sender));
 
         // Return the subscription ID and entity change stream
         StoreEventStream::new(Box::new(ReceiverStream::new(receiver).map(Ok).compat()))
             .filter_by_entities(entities)
     }
 
-    fn subscribe_no_payload(&self, entities: Vec<SubscriptionFilter>) -> UnitStream {
-        let id = Uuid::new_v4().to_string();
-
-        let (sender, receiver) = watch::channel(());
-
-        self.subscriptions
-            .write()
-            .unwrap()
-            .insert(id, (Arc::new(entities.clone()), Arc::new(sender)));
-
-        Box::new(tokio_stream::wrappers::WatchStream::new(receiver))
+    fn subscribe_no_payload(&self, entities: BTreeSet<SubscriptionFilter>) -> UnitStream {
+        self.subscriptions_no_payload
+            .lock()
+            .entry(entities)
+            .or_insert_with(|| Watcher::new(()))
+            .stream()
     }
 }

store/postgres/tests/store.rs

@@ -975,10 +975,11 @@ fn subscribe(
     subgraph: &DeploymentHash,
     entity_type: &str,
 ) -> StoreEventStream<impl Stream<Item = Arc<StoreEvent>, Error = ()> + Send> {
-    let subscription = SUBSCRIPTION_MANAGER.subscribe(vec![SubscriptionFilter::Entities(
-        subgraph.clone(),
-        EntityType::new(entity_type.to_owned()),
-    )]);
+    let subscription =
+        SUBSCRIPTION_MANAGER.subscribe(FromIterator::from_iter([SubscriptionFilter::Entities(
+            subgraph.clone(),
+            EntityType::new(entity_type.to_owned()),
+        )]));
 
     StoreEventStream::new(subscription)
 }