ethereum, store: Replace chain head listener channels with tokio::watch

Simplifies and improves performance of the
communication between the listener and block streams.

Author: leoyvens

datasource/ethereum/src/block_stream.rs

@@ -7,7 +7,6 @@ use std::sync::Mutex;
 use std::time::{Duration, Instant};
 
 use futures::prelude::*;
-use futures::sync::mpsc::{channel, Receiver, Sender};
 use graph::tokio::timer::Delay;
 
 use graph::data::subgraph::schema::{
@@ -152,9 +151,7 @@ impl<S, C, E> Clone for BlockStreamContext<S, C, E> {
 pub struct BlockStream<S, C, E> {
     state: Mutex<BlockStreamState>,
     consecutive_err_count: u32,
-    chain_head_update_sink: Sender<ChainHeadUpdate>,
-    chain_head_update_stream: Receiver<ChainHeadUpdate>,
-    _chain_head_update_guard: CancelGuard,
+    chain_head_update_stream: ChainHeadUpdateStream,
     ctx: BlockStreamContext<S, C, E>,
 }
 
@@ -168,7 +165,6 @@ where
         subgraph_store: Arc<S>,
         chain_store: Arc<C>,
         eth_adapter: Arc<E>,
-        chain_head_update_guard: CancelGuard,
         node_id: NodeId,
         subgraph_id: SubgraphDeploymentId,
         log_filter: Option<EthereumLogFilter>,
@@ -178,14 +174,10 @@ where
         reorg_threshold: u64,
         logger: Logger,
     ) -> Self {
-        let (chain_head_update_sink, chain_head_update_stream) = channel(100);
-
         BlockStream {
             state: Mutex::new(BlockStreamState::New),
             consecutive_err_count: 0,
-            chain_head_update_sink,
-            chain_head_update_stream,
-            _chain_head_update_guard: chain_head_update_guard,
+            chain_head_update_stream: chain_store.chain_head_updates(),
             ctx: BlockStreamContext {
                 subgraph_store,
                 chain_store,
@@ -1181,7 +1173,7 @@ where
                 BlockStreamState::Idle => {
                     match self.chain_head_update_stream.poll() {
                         // Chain head was updated
-                        Ok(Async::Ready(Some(_chain_head_update))) => {
+                        Ok(Async::Ready(Some(()))) => {
                             // Start reconciliation process
                             let next_blocks_future = self.ctx.next_blocks();
                             state = BlockStreamState::Reconciliation(next_blocks_future);
@@ -1222,21 +1214,6 @@ where
     }
 }
 
-impl<S, C, E> EventConsumer<ChainHeadUpdate> for BlockStream<S, C, E>
-where
-    S: Store,
-    C: ChainStore,
-    E: EthereumAdapter,
-{
-    fn event_sink(&self) -> Box<Sink<SinkItem = ChainHeadUpdate, SinkError = ()> + Send> {
-        let logger = self.ctx.logger.clone();
-
-        Box::new(self.chain_head_update_sink.clone().sink_map_err(move |_| {
-            debug!(logger, "Terminating chain head updates; channel closed");
-        }))
-    }
-}
-
 pub struct BlockStreamBuilder<S, C, E> {
     subgraph_store: Arc<S>,
     chain_store: Arc<C>,
@@ -1300,25 +1277,12 @@ where
         let logger = logger.new(o!(
             "component" => "BlockStream",
         ));
-        let logger_for_stream = logger.clone();
-
-        // Create a chain head update stream whose lifetime is tied to the
-        // liftetime of the block stream; we do this to immediately terminate
-        // the chain head update listener when the block stream is shut down
-        let cancel_guard = CancelGuard::new();
-        let chain_head_update_stream =
-            self.chain_store
-                .chain_head_updates()
-                .cancelable(&cancel_guard, move || {
-                    debug!(logger_for_stream, "Terminating chain head updates");
-                });
 
         // Create the actual subgraph-specific block stream
-        let block_stream = BlockStream::new(
+        BlockStream::new(
             self.subgraph_store.clone(),
             self.chain_store.clone(),
             self.eth_adapter.clone(),
-            cancel_guard,
             self.node_id.clone(),
             deployment_id,
             log_filter,
@@ -1327,18 +1291,7 @@ where
             include_calls_in_blocks,
             self.reorg_threshold,
             logger,
-        );
-
-        // Forward chain head updates from the listener to the block stream;
-        // this will be canceled as soon as the block stream goes out of scope
-        tokio::spawn(
-            chain_head_update_stream
-                .forward(block_stream.event_sink())
-                .map_err(|_| ())
-                .map(|_| ()),
-        );
-
-        block_stream
+        )
     }
 }
 

graph/src/components/ethereum/listener.rs

@@ -21,7 +21,9 @@ pub struct ChainHeadUpdate {
     pub head_block_number: u64,
 }
 
-pub type ChainHeadUpdateStream = Box<Stream<Item = ChainHeadUpdate, Error = ()> + Send>;
+/// The updates have no payload, receivers should call `Store::chain_head_ptr`
+/// to check what the latest block is.
+pub type ChainHeadUpdateStream = Box<Stream<Item = (), Error = ()> + Send>;
 
 pub trait ChainHeadUpdateListener {
     // Subscribe to chain head updates.

graph/src/components/ethereum/stream.rs

@@ -3,9 +3,7 @@ use futures::Stream;
 
 use crate::prelude::*;
 
-pub trait BlockStream:
-    Stream<Item = EthereumBlockWithTriggers, Error = Error> + EventConsumer<ChainHeadUpdate>
-{
+pub trait BlockStream: Stream<Item = EthereumBlockWithTriggers, Error = Error> {
     fn parse_triggers(
         log_filter_opt: Option<EthereumLogFilter>,
         call_filter_opt: Option<EthereumCallFilter>,

store/postgres/src/chain_head_listener.rs

@@ -1,25 +1,18 @@
-use futures::sync::mpsc::{channel, Sender};
-use std::collections::HashMap;
-use std::sync::{Arc, RwLock};
-use uuid::Uuid;
+use tokio::sync::watch;
 
 use graph::prelude::{ChainHeadUpdateListener as ChainHeadUpdateListenerTrait, *};
 use graph::serde_json;
 
 use crate::notification_listener::{NotificationListener, SafeChannelName};
 
-type ChainHeadUpdateSubscribers = Arc<RwLock<HashMap<String, Sender<ChainHeadUpdate>>>>;
-
 pub struct ChainHeadUpdateListener {
-    logger: Logger,
-    subscribers: ChainHeadUpdateSubscribers,
+    update_receiver: watch::Receiver<()>,
     _listener: NotificationListener,
 }
 
 impl ChainHeadUpdateListener {
     pub fn new(logger: &Logger, postgres_url: String, network_name: String) -> Self {
         let logger = logger.new(o!("component" => "ChainHeadUpdateListener"));
-        let subscribers = Arc::new(RwLock::new(HashMap::new()));
 
         // Create a Postgres notification listener for chain head updates
         let mut listener = NotificationListener::new(
@@ -28,11 +21,11 @@ impl ChainHeadUpdateListener {
             SafeChannelName::i_promise_this_is_safe("chain_head_updates"),
         );
 
-        Self::listen(&logger, &mut listener, network_name, subscribers.clone());
+        let (update_sender, update_receiver) = watch::channel(());
+        Self::listen(logger, &mut listener, network_name, update_sender);
 
         ChainHeadUpdateListener {
-            logger,
-            subscribers,
+            update_receiver,
 
             // We keep the listener around to tie its stream's lifetime to
             // that of the chain head update listener and prevent it from
@@ -42,10 +35,10 @@ impl ChainHeadUpdateListener {
     }
 
     fn listen(
-        logger: &Logger,
+        logger: Logger,
         listener: &mut NotificationListener,
         network_name: String,
-        subscribers: ChainHeadUpdateSubscribers,
+        mut update_sender: watch::Sender<()>,
     ) {
         let logger = logger.clone();
 
@@ -72,64 +65,25 @@ impl ChainHeadUpdateListener {
                     }
                 })
                 .for_each(move |update| {
-                    let logger = logger.clone();
-                    let senders = subscribers.read().unwrap().clone();
-                    let subscribers = subscribers.clone();
-
                     debug!(
-                        logger,
+                        logger.clone(),
                         "Received chain head update";
                         "network" => &update.network_name,
                         "head_block_hash" => format!("{}", update.head_block_hash),
                         "head_block_number" => &update.head_block_number,
                     );
 
-                    // Forward update to all susbcribers
-                    stream::iter_ok::<_, ()>(senders).for_each(move |(id, mut sender)| {
-                        let logger = logger.clone();
-                        let subscribers = subscribers.clone();
-
-                        // A subgraph that's syncing will let chain head updates
-                        // pile up in the channel. So we don't wait for room in
-                        // the channel and instead skip it, it will have the
-                        // opportunity to grab future updates.
-                        match sender.try_send(update.clone()) {
-                            // Move on to the next subscriber
-                            Ok(()) => (),
-                            Err(ref e) if e.is_full() => {
-                                // Temporary log, feel free to remove if noisy.
-                                debug!(logger, "Full chain head update channel"; "id" => &id);
-                            }
-                            Err(ref e) if e.is_disconnected() => {
-                                // Remove disconnected subscribers.
-                                debug!(logger, "Unsubscribe"; "id" => &id);
-                                subscribers.write().unwrap().remove(&id);
-                            }
-                            Err(e) => warn!(logger, "Unexpected send error"; "e" => e.to_string()),
-                        }
-                        Ok(())
-                    })
+                    update_sender.broadcast(()).map_err(|_| ())
                 }),
         );
 
-        // We're ready, start listening to chain head updaates
+        // We're ready, start listening to chain head updates
         listener.start();
     }
 }
 
 impl ChainHeadUpdateListenerTrait for ChainHeadUpdateListener {
     fn subscribe(&self) -> ChainHeadUpdateStream {
-        // Generate a new (unique) UUID; we're looping just to be sure we avoid collisions
-        let mut id = Uuid::new_v4().to_string();
-        while self.subscribers.read().unwrap().contains_key(&id) {
-            id = Uuid::new_v4().to_string();
-        }
-
-        debug!(self.logger, "Subscribe"; "id" => &id);
-
-        // Create a subscriber and return the receiving end
-        let (sender, receiver) = channel(100);
-        self.subscribers.write().unwrap().insert(id, sender);
-        Box::new(receiver)
+        Box::new(self.update_receiver.clone().map_err(|_| ()))
     }
 }