chain head listener: change watcher deadlock detection to avoid spurious deadlock detection

leoyvens · leoyvens · commit c28c23355db2 · 2021-12-15T11:10:45.000Z
diff --git a/store/postgres/src/chain_head_listener.rs b/store/postgres/src/chain_head_listener.rs
@@ -2,14 +2,15 @@ use graph::{
     blockchain::ChainHeadUpdateStream,
     prelude::{
         futures03::{self, FutureExt},
-        tokio, StoreError,
+        tokio::{self, task::JoinHandle},
+        StoreError,
     },
     prometheus::{CounterVec, GaugeVec},
     util::timed_rw_lock::TimedRwLock,
 };
-use std::sync::{atomic::AtomicBool, Arc};
+use std::str::FromStr;
+use std::sync::Arc;
 use std::{collections::BTreeMap, time::Duration};
-use std::{str::FromStr, sync::atomic};
 
 use lazy_static::lazy_static;
 
@@ -151,7 +152,7 @@ impl ChainHeadUpdateListener {
     ) {
         // Process chain head updates in a dedicated task
         graph::spawn(async move {
-            let sending_to_watcher = Arc::new(AtomicBool::new(false));
+            let mut send_to_watcher: Option<JoinHandle<()>> = None;
             while let Some(notification) = receiver.recv().await {
                 // Create ChainHeadUpdate from JSON
                 let update: ChainHeadUpdate =
@@ -180,16 +181,16 @@ impl ChainHeadUpdateListener {
                 if let Some(watcher) = watchers.read(&logger).get(&update.network_name) {
                     // Due to a tokio bug, we must assume that the watcher can deadlock, see
                     // https://github.com/tokio-rs/tokio/issues/4246.
-                    if !sending_to_watcher.load(atomic::Ordering::SeqCst) {
-                        let sending_to_watcher = sending_to_watcher.cheap_clone();
-                        let sender = watcher.sender.cheap_clone();
-                        tokio::task::spawn_blocking(move || {
-                            sending_to_watcher.store(true, atomic::Ordering::SeqCst);
-                            sender.send(()).unwrap();
-                            sending_to_watcher.store(false, atomic::Ordering::SeqCst);
-                        });
-                    } else {
-                        debug!(logger, "skipping chain head update, watcher is deadlocked"; "network" => &update.network_name);
+                    match send_to_watcher.as_mut().map(FutureExt::now_or_never) {
+                        None | Some(Some(_)) => {
+                            let sender = watcher.sender.cheap_clone();
+                            send_to_watcher = Some(tokio::task::spawn_blocking(move || {
+                                sender.send(()).unwrap()
+                            }));
+                        }
+                        Some(None) => {
+                            debug!(logger, "skipping chain head update, watcher is deadlocked"; "network" => &update.network_name)
+                        }
                     }
                 }
             }
