pollinator: Added Receiver struct

Author: DougLau

pollinator/src/bin/pollinator.rs

@@ -65,15 +65,15 @@ struct CommLinkTask {
 
 /// Poll comm links
 async fn poll_comm_links(db: Database) -> Result<()> {
-    let (notifier, mut stream) = db
-        .clone()
-        .notifier(
-            ["comm_config", "comm_link", "controller", "detector"].into_iter(),
-        )
-        .await?;
+    let (notifier, mut receiver) = db.clone().notifier().await?;
     let mut tasks: HashMap<String, CommLinkTask> = HashMap::new();
     let mut set = JoinSet::new();
     set.spawn(notifier.run());
+    receiver
+        .listen(
+            ["comm_config", "comm_link", "controller", "detector"].into_iter(),
+        )
+        .await?;
     loop {
         let cfgs = CommLinkCfg::lookup_all(db.clone()).await?;
         if cfgs.is_empty() {
@@ -112,17 +112,17 @@ async fn poll_comm_links(db: Database) -> Result<()> {
                 tasks.insert(name, task);
             }
         }
-        while let Some(not) = stream.recv().await {
+        while let Some(not) = receiver.recv().await {
             if should_reload(&not, &cfgs) {
                 log::info!(
                     "{} {}: reloading configuration",
                     not.channel(),
                     not.payload()
                 );
                 tokio::time::sleep(Duration::from_secs(2)).await;
-                // Empty the notification stream
-                while !stream.is_empty() {
-                    stream.recv().await;
+                // Empty the notification receiver
+                while !receiver.is_empty() {
+                    receiver.recv().await;
                 }
                 break;
             }

resin/src/notify.rs

@@ -24,15 +24,20 @@ use tokio_postgres::{AsyncMessage, Client, Connection, Notification, Socket};
 
 /// Handler for Postgres NOTIFY
 pub struct Notifier {
-    /// DB client (kept so it's not dropped)
-    #[allow(unused)]
-    client: Client,
     /// DB connection
     conn: Connection<Socket, NoTlsStream>,
     /// Notification sender
     tx: UnboundedSender<Notification>,
 }
 
+/// Receiver for Postgres NOTIFY
+pub struct Receiver {
+    /// DB client
+    client: Client,
+    /// Notification receiver
+    rx: UnboundedReceiver<Notification>,
+}
+
 impl Future for Notifier {
     type Output = bool;
 
@@ -72,18 +77,38 @@ impl Notifier {
     }
 }
 
-impl Database {
-    /// Create a new notifier
-    pub async fn notifier(
-        self,
+impl Receiver {
+    /// Listen for notifications on channels
+    pub async fn listen(
+        &self,
         channels: impl Iterator<Item = &str>,
-    ) -> Result<(Notifier, UnboundedReceiver<Notification>)> {
-        let (client, conn) = self.dedicated_client().await?;
-        let (tx, rx) = unbounded_channel();
+    ) -> Result<()> {
         for channel in channels {
-            client.execute(&format!("LISTEN {channel}"), &[]).await?;
+            self.client
+                .execute(&format!("LISTEN {channel}"), &[])
+                .await?;
         }
-        let not = Notifier { client, conn, tx };
-        Ok((not, rx))
+        Ok(())
+    }
+
+    /// Receive one notification
+    pub async fn recv(&mut self) -> Option<Notification> {
+        self.rx.recv().await
+    }
+
+    /// Check if channel is empty
+    pub fn is_empty(&self) -> bool {
+        self.rx.is_empty()
+    }
+}
+
+impl Database {
+    /// Create a new notifier / receiver
+    pub async fn notifier(self) -> Result<(Notifier, Receiver)> {
+        let (client, conn) = self.dedicated_client().await?;
+        let (tx, rx) = unbounded_channel();
+        let not = Notifier { conn, tx };
+        let rcv = Receiver { client, rx };
+        Ok((not, rcv))
     }
 }