Send publisher permissions using a watch channel (#118)

Fixes heap growth if Exporter::update_our_prices is not called as often as updates

Author: alexheretic

src/agent/solana.rs

@@ -8,9 +8,9 @@ pub mod network {
 
     use {
         super::{
-            super::{
-                store,
-                store::global,
+            super::store::{
+                self,
+                global,
             },
             exporter,
             key_store::{
@@ -29,8 +29,8 @@ pub mod network {
         std::time::Duration,
         tokio::{
             sync::{
-                mpsc,
                 mpsc::Sender,
+                watch,
             },
             task::JoinHandle,
         },
@@ -86,8 +86,7 @@ pub mod network {
         logger: Logger,
     ) -> Result<Vec<JoinHandle<()>>> {
         // Publisher permissions updates between oracle and exporter
-        let (publisher_permissions_tx, publisher_permissions_rx) =
-            mpsc::channel(config.oracle.updates_channel_capacity);
+        let (publisher_permissions_tx, publisher_permissions_rx) = watch::channel(<_>::default());
 
         // Spawn the Oracle
         let mut jhs = oracle::spawn_oracle(

src/agent/solana/exporter.rs

@@ -172,7 +172,7 @@ pub fn spawn_exporter(
     network: Network,
     rpc_url: &str,
     rpc_timeout: Duration,
-    publisher_permissions_rx: mpsc::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+    publisher_permissions_rx: watch::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
     key_store: KeyStore,
     local_store_tx: Sender<store::local::Message>,
     global_store_tx: Sender<store::global::Lookup>,
@@ -260,7 +260,7 @@ pub struct Exporter {
     inflight_transactions_tx: Sender<Signature>,
 
     /// publisher => { permissioned_price => market hours } as read by the oracle module
-    publisher_permissions_rx: mpsc::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+    publisher_permissions_rx: watch::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
 
     /// Currently known permissioned prices of this publisher along with their market hours
     our_prices: HashMap<Pubkey, MarketSchedule>,
@@ -287,7 +287,7 @@ impl Exporter {
         global_store_tx: Sender<store::global::Lookup>,
         network_state_rx: watch::Receiver<NetworkState>,
         inflight_transactions_tx: Sender<Signature>,
-        publisher_permissions_rx: mpsc::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+        publisher_permissions_rx: watch::Receiver<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
         keypair_request_tx: mpsc::Sender<KeypairRequest>,
         logger: Logger,
     ) -> Self {
@@ -556,51 +556,34 @@ impl Exporter {
     /// The loop ensures that we clear the channel and use
     /// only the final, latest message; try_recv() is
     /// non-blocking.
-    ///
-    /// Note: This behavior is similar to
-    /// tokio::sync::watch::channel(), which was not appropriate here
-    /// because its internal RwLock would complain about not being
-    /// Send with the HashMap<HashSet<Pubkey>> inside.
-    /// TODO(2023-05-05): Debug the watch::channel() compilation errors
     fn update_our_prices(&mut self, publish_pubkey: &Pubkey) {
-        loop {
-            match self.publisher_permissions_rx.try_recv() {
-                Ok(publisher_permissions) => {
-                    self.our_prices = publisher_permissions.get(publish_pubkey) .cloned()
-                        .unwrap_or_else( || {
-                            warn!(
-                                self.logger,
-                                "Exporter: No permissioned prices were found for the publishing keypair on-chain. This is expected only on startup.";
-                                "publish_pubkey" => publish_pubkey.to_string(),
-                            );
-                            HashMap::new()
-                        });
-                    trace!(
-                        self.logger,
-                        "Exporter: read permissioned price accounts from channel";
-                        "new_value" => format!("{:?}", self.our_prices),
-                    );
-                }
-                // Expected failures when channel is empty
-                Err(TryRecvError::Empty) => {
-                    trace!(
-                        self.logger,
-                        "Exporter: No more permissioned price accounts in channel, using cached value";
-                    );
-                    break;
-                }
-                // Unexpected failures (channel closed, internal errors etc.)
-                Err(other) => {
-                    warn!(
-                        self.logger,
-                        "Exporter: Updating permissioned price accounts failed unexpectedly, using cached value";
-                        "cached_value" => format!("{:?}", self.our_prices),
-                        "error" => other.to_string(),
-                    );
-                    break;
-                }
+        match self.publisher_permissions_rx.has_changed() {
+            Ok(true) => {}
+            Ok(false) => return,
+            Err(other) => {
+                warn!(
+                    self.logger,
+                    "Exporter: Updating permissioned price accounts failed unexpectedly, using cached value";
+                    "cached_value" => format!("{:?}", self.our_prices),
+                    "error" => other.to_string(),
+                );
+                return;
             }
         }
+
+        self.our_prices = self
+            .publisher_permissions_rx
+            .borrow_and_update()
+            .get(publish_pubkey)
+            .cloned()
+            .unwrap_or_else(|| {
+                warn!(
+                    self.logger,
+                    "Exporter: No permissioned prices were found for the publishing keypair on-chain. This is expected only on startup.";
+                    "publish_pubkey" => publish_pubkey.to_string(),
+                );
+                HashMap::new()
+            });
     }
 
     async fn fetch_local_store_contents(&self) -> Result<HashMap<PriceIdentifier, PriceInfo>> {

src/agent/solana/oracle.rs

@@ -43,7 +43,10 @@ use {
         time::Duration,
     },
     tokio::{
-        sync::mpsc,
+        sync::{
+            mpsc,
+            watch,
+        },
         task::JoinHandle,
         time::Interval,
     },
@@ -204,7 +207,7 @@ pub fn spawn_oracle(
     wss_url: &str,
     rpc_timeout: Duration,
     global_store_update_tx: mpsc::Sender<global::Update>,
-    publisher_permissions_tx: mpsc::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+    publisher_permissions_tx: watch::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
     key_store: KeyStore,
     logger: Logger,
 ) -> Vec<JoinHandle<()>> {
@@ -419,7 +422,7 @@ struct Poller {
     data_tx: mpsc::Sender<Data>,
 
     /// Updates about permissioned price accounts from oracle to exporter
-    publisher_permissions_tx: mpsc::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+    publisher_permissions_tx: watch::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
 
     /// The RPC client to use to poll data from the RPC node
     rpc_client: RpcClient,
@@ -439,7 +442,7 @@ struct Poller {
 impl Poller {
     pub fn new(
         data_tx: mpsc::Sender<Data>,
-        publisher_permissions_tx: mpsc::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
+        publisher_permissions_tx: watch::Sender<HashMap<Pubkey, HashMap<Pubkey, MarketSchedule>>>,
         rpc_url: &str,
         rpc_timeout: Duration,
         commitment: CommitmentLevel,
@@ -481,9 +484,7 @@ impl Poller {
         let fresh_data = self.poll().await?;
 
         self.publisher_permissions_tx
-            .send(fresh_data.publisher_permissions.clone())
-            .await
-            .context("Updating permissioned price accounts for exporter")?;
+            .send_replace(fresh_data.publisher_permissions.clone());
 
         self.data_tx
             .send(fresh_data)