refactor: cleaner sender_denylist_watcher shutdown

Signed-off-by: Alexis Asseman <[email protected]>
graphprotocol · Jan 10, 2024 · 6f263c5 · 6f263c5
1 parent a2e4621
commit 6f263c5
Show file tree

Hide file tree

Showing 3 changed files with 58 additions and 39 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/common/Cargo.toml b/common/Cargo.toml
@@ -45,6 +45,7 @@ autometrics = { version = "0.6.0", features = ["prometheus-exporter"] }
 tracing = "0.1.40"
 tower = "0.4.13"
 tower_governor = "0.1.0"
+tokio-util = "0.7.10"
 
 [dev-dependencies]
 env_logger = "0.9.0"

diff --git a/common/src/tap_manager.rs b/common/src/tap_manager.rs
@@ -23,6 +23,7 @@ pub struct TapManager {
     domain_separator: Arc<Eip712Domain>,
     sender_denylist: Arc<RwLock<HashSet<Address>>>,
     _sender_denylist_watcher_handle: Arc<tokio::task::JoinHandle<()>>,
+    sender_denylist_watcher_cancel_token: Arc<tokio_util::sync::CancellationToken>,
 }
 
 impl TapManager {
@@ -49,10 +50,13 @@ impl TapManager {
             .await
             .expect("should be able to fetch the sender_denylist from the DB on startup");
 
+        let sender_denylist_watcher_cancel_token =
+            Arc::new(tokio_util::sync::CancellationToken::new());
         let sender_denylist_watcher_handle = Arc::new(tokio::spawn(Self::sender_denylist_watcher(
             pgpool.clone(),
             pglistener,
             sender_denylist.clone(),
+            sender_denylist_watcher_cancel_token.clone(),
         )));
 
         Self {
@@ -62,6 +66,7 @@ impl TapManager {
             domain_separator: Arc::new(domain_separator),
             sender_denylist,
             _sender_denylist_watcher_handle: sender_denylist_watcher_handle,
+            sender_denylist_watcher_cancel_token,
         }
     }
 
@@ -169,6 +174,7 @@ impl TapManager {
         pgpool: PgPool,
         mut pglistener: PgListener,
         denylist: Arc<RwLock<HashSet<Address>>>,
+        cancel_token: Arc<tokio_util::sync::CancellationToken>,
     ) {
         #[derive(serde::Deserialize)]
         struct DenylistNotification {
@@ -177,52 +183,63 @@ impl TapManager {
         }
 
         loop {
-            let pg_notification = pglistener.recv().await.expect(
-                "should be able to receive Postgres Notify events on the channel \
-                'scalar_tap_deny_notification'",
-            );
-
-            println!(
-                "Received a denylist table notification: {}",
-                pg_notification.payload()
-            );
-
-            let denylist_notification: DenylistNotification =
-                serde_json::from_str(pg_notification.payload()).expect(
-                    "should be able to deserialize the Postgres Notify event payload as a \
-                    DenylistNotification",
-                );
-
-            match denylist_notification.tg_op.as_str() {
-                "INSERT" => {
-                    denylist
-                        .write()
-                        .await
-                        .insert(denylist_notification.sender_address);
+            tokio::select! {
+                _ = cancel_token.cancelled() => {
+                    break;
                 }
-                "DELETE" => {
-                    denylist
-                        .write()
-                        .await
-                        .remove(&denylist_notification.sender_address);
-                }
-                // UPDATE and TRUNCATE are not expected to happen. Reload the entire denylist.
-                _ => {
-                    error!(
-                        "Received an unexpected denylist table notification: {}. Reloading entire \
-                        denylist.",
-                        denylist_notification.tg_op
+
+                pg_notification = pglistener.recv() => {
+                    let pg_notification = pg_notification.expect(
+                    "should be able to receive Postgres Notify events on the channel \
+                    'scalar_tap_deny_notification'",
                     );
 
-                    Self::sender_denylist_reload(pgpool.clone(), denylist.clone())
-                        .await
-                        .expect("should be able to reload the sender denylist")
+                    let denylist_notification: DenylistNotification =
+                        serde_json::from_str(pg_notification.payload()).expect(
+                            "should be able to deserialize the Postgres Notify event payload as a \
+                            DenylistNotification",
+                        );
+
+                    match denylist_notification.tg_op.as_str() {
+                        "INSERT" => {
+                            denylist
+                                .write()
+                                .await
+                                .insert(denylist_notification.sender_address);
+                        }
+                        "DELETE" => {
+                            denylist
+                                .write()
+                                .await
+                                .remove(&denylist_notification.sender_address);
+                        }
+                        // UPDATE and TRUNCATE are not expected to happen. Reload the entire denylist.
+                        _ => {
+                            error!(
+                                "Received an unexpected denylist table notification: {}. Reloading entire \
+                                denylist.",
+                                denylist_notification.tg_op
+                            );
+
+                            Self::sender_denylist_reload(pgpool.clone(), denylist.clone())
+                                .await
+                                .expect("should be able to reload the sender denylist")
+                        }
+                    }
                 }
             }
         }
     }
 }
 
+impl Drop for TapManager {
+    fn drop(&mut self) {
+        // Clean shutdown for the sender_denylist_watcher
+        // Though since it's not a critical task, we don't wait for it to finish (join).
+        self.sender_denylist_watcher_cancel_token.cancel();
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::str::FromStr;