fix: remove duplicate code

Author: SupernaviX

common/src/caryatid.rs

@@ -0,0 +1,66 @@
+use std::sync::Arc;
+
+use anyhow::Result;
+use caryatid_sdk::{async_trait, Context, MessageBounds, Subscription};
+
+use crate::messages::{CardanoMessage, Message};
+
+#[async_trait]
+pub trait SubscriptionExt<M: MessageBounds> {
+    async fn read_ignoring_rollbacks(&mut self) -> Result<(String, Arc<M>)>;
+}
+
+#[async_trait]
+impl SubscriptionExt<Message> for Box<dyn Subscription<Message>> {
+    async fn read_ignoring_rollbacks(&mut self) -> Result<(String, Arc<Message>)> {
+        loop {
+            let (stream, message) = self.read().await?;
+            if matches!(
+                message.as_ref(),
+                Message::Cardano((_, CardanoMessage::Rollback(_)))
+            ) {
+                continue;
+            }
+            break Ok((stream, message));
+        }
+    }
+}
+
+/// A utility to publish messages, which will only publish rollback messages if some work has been rolled back
+pub struct RollbackAwarePublisher<M: MessageBounds> {
+    /// Module context
+    context: Arc<Context<M>>,
+
+    /// Topic to publish on
+    topic: String,
+
+    // When did we publish our last non-rollback message
+    last_activity_at: Option<u64>,
+}
+
+impl RollbackAwarePublisher<Message> {
+    pub fn new(context: Arc<Context<Message>>, topic: String) -> Self {
+        Self {
+            context,
+            topic,
+            last_activity_at: None,
+        }
+    }
+
+    pub async fn publish(&mut self, message: Arc<Message>) -> Result<()> {
+        match message.as_ref() {
+            Message::Cardano((block, CardanoMessage::Rollback(_))) => {
+                if self.last_activity_at.is_some_and(|slot| slot >= block.slot) {
+                    self.last_activity_at = None;
+                    self.context.publish(&self.topic, message).await?;
+                }
+                Ok(())
+            }
+            Message::Cardano((block, _)) => {
+                self.last_activity_at = Some(block.slot);
+                self.context.publish(&self.topic, message).await
+            }
+            _ => self.context.publish(&self.topic, message).await,
+        }
+    }
+}

common/src/lib.rs

@@ -2,6 +2,7 @@
 
 pub mod address;
 pub mod calculations;
+pub mod caryatid;
 pub mod cbor;
 pub mod cip19;
 pub mod commands;
@@ -22,7 +23,6 @@ pub mod serialization;
 pub mod snapshot;
 pub mod stake_addresses;
 pub mod state_history;
-pub mod subscription;
 pub mod types;
 pub mod upstream_cache;
 pub mod validation;

common/src/subscription.rs

@@ -2,10 +2,10 @@
 //! Manages stake and reward accounts state
 
 use acropolis_common::{
+    caryatid::SubscriptionExt,
     messages::{CardanoMessage, Message, StateQuery, StateQueryResponse},
     queries::accounts::{DrepDelegators, PoolDelegators, DEFAULT_ACCOUNTS_QUERY_TOPIC},
     state_history::{StateHistory, StateHistoryStore},
-    subscription::SubscriptionExt,
     BlockInfo, BlockStatus,
 };
 use anyhow::Result;

modules/accounts_state/src/accounts_state.rs

@@ -1,3 +1,4 @@
+use acropolis_common::caryatid::RollbackAwarePublisher;
 use acropolis_common::messages::{
     CardanoMessage, DRepDelegationDistribution, DRepStakeDistributionMessage, Message,
 };
@@ -6,25 +7,12 @@ use caryatid_sdk::Context;
 use std::sync::Arc;
 
 /// Message publisher for DRep Delegation Distribution (DRDD)
-pub struct DRepDistributionPublisher {
-    /// Module context
-    context: Arc<Context<Message>>,
-
-    /// Topic to publish on
-    topic: String,
-
-    // When did we publish our last non-rollback message
-    last_activity_at: Option<u64>,
-}
+pub struct DRepDistributionPublisher(RollbackAwarePublisher<Message>);
 
 impl DRepDistributionPublisher {
     /// Construct with context and topic to publish on
     pub fn new(context: Arc<Context<Message>>, topic: String) -> Self {
-        Self {
-            context,
-            topic,
-            last_activity_at: None,
-        }
+        Self(RollbackAwarePublisher::new(context, topic))
     }
 
     /// Publish the DRep Delegation Distribution
@@ -33,31 +21,19 @@ impl DRepDistributionPublisher {
         block: &BlockInfo,
         drdd: DRepDelegationDistribution,
     ) -> anyhow::Result<()> {
-        self.last_activity_at = Some(block.slot);
-        self.context
-            .message_bus
-            .publish(
-                &self.topic,
-                Arc::new(Message::Cardano((
-                    block.clone(),
-                    CardanoMessage::DRepStakeDistribution(DRepStakeDistributionMessage {
-                        epoch: block.epoch,
-                        drdd,
-                    }),
-                ))),
-            )
+        self.0
+            .publish(Arc::new(Message::Cardano((
+                block.clone(),
+                CardanoMessage::DRepStakeDistribution(DRepStakeDistributionMessage {
+                    epoch: block.epoch,
+                    drdd,
+                }),
+            ))))
             .await
     }
 
     /// Publish a rollback message, if we have anything to roll back
     pub async fn publish_rollback(&mut self, message: Arc<Message>) -> anyhow::Result<()> {
-        let Message::Cardano((block_info, CardanoMessage::Rollback(_))) = message.as_ref() else {
-            return Ok(());
-        };
-        if self.last_activity_at.is_none_or(|slot| slot < block_info.slot) {
-            return Ok(());
-        }
-        self.last_activity_at = None;
-        self.context.message_bus.publish(&self.topic, message).await
+        self.0.publish(message).await
     }
 }

modules/accounts_state/src/drep_distribution_publisher.rs

@@ -1,29 +1,17 @@
+use acropolis_common::caryatid::RollbackAwarePublisher;
 use acropolis_common::messages::{CardanoMessage, Message, SPOStakeDistributionMessage};
 use acropolis_common::{BlockInfo, DelegatedStake, PoolId};
 use caryatid_sdk::Context;
 use std::collections::BTreeMap;
 use std::sync::Arc;
 
 /// Message publisher for Stake Pool Delegation Distribution (SPDD)
-pub struct SPODistributionPublisher {
-    /// Module context
-    context: Arc<Context<Message>>,
-
-    /// Topic to publish on
-    topic: String,
-
-    // When did we publish our last non-rollback message
-    last_activity_at: Option<u64>,
-}
+pub struct SPODistributionPublisher(RollbackAwarePublisher<Message>);
 
 impl SPODistributionPublisher {
     /// Construct with context and topic to publish on
     pub fn new(context: Arc<Context<Message>>, topic: String) -> Self {
-        Self {
-            context,
-            topic,
-            last_activity_at: None,
-        }
+        Self(RollbackAwarePublisher::new(context, topic))
     }
 
     /// Publish the SPDD
@@ -32,31 +20,19 @@ impl SPODistributionPublisher {
         block: &BlockInfo,
         spos: BTreeMap<PoolId, DelegatedStake>,
     ) -> anyhow::Result<()> {
-        self.last_activity_at = Some(block.slot);
-        self.context
-            .message_bus
-            .publish(
-                &self.topic,
-                Arc::new(Message::Cardano((
-                    block.clone(),
-                    CardanoMessage::SPOStakeDistribution(SPOStakeDistributionMessage {
-                        epoch: block.epoch - 1, // End of the previous epoch
-                        spos: spos.into_iter().collect(),
-                    }),
-                ))),
-            )
+        self.0
+            .publish(Arc::new(Message::Cardano((
+                block.clone(),
+                CardanoMessage::SPOStakeDistribution(SPOStakeDistributionMessage {
+                    epoch: block.epoch - 1, // End of the previous epoch
+                    spos: spos.into_iter().collect(),
+                }),
+            ))))
             .await
     }
 
     /// Publish a rollback message, if we have anything to roll back
     pub async fn publish_rollback(&mut self, message: Arc<Message>) -> anyhow::Result<()> {
-        let Message::Cardano((block_info, CardanoMessage::Rollback(_))) = message.as_ref() else {
-            return Ok(());
-        };
-        if self.last_activity_at.is_none_or(|slot| slot < block_info.slot) {
-            return Ok(());
-        }
-        self.last_activity_at = None;
-        self.context.message_bus.publish(&self.topic, message).await
+        self.0.publish(message).await
     }
 }

modules/accounts_state/src/spo_distribution_publisher.rs

@@ -1,28 +1,16 @@
+use acropolis_common::caryatid::RollbackAwarePublisher;
 use acropolis_common::messages::{CardanoMessage, Message, SPORewardsMessage};
 use acropolis_common::{BlockInfo, PoolId, SPORewards};
 use caryatid_sdk::Context;
 use std::sync::Arc;
 
 /// Message publisher for Stake Pool Delegation Distribution (SPDD)
-pub struct SPORewardsPublisher {
-    /// Module context
-    context: Arc<Context<Message>>,
-
-    /// Topic to publish on
-    topic: String,
-
-    // When did we publish our last non-rollback message
-    last_activity_at: Option<u64>,
-}
+pub struct SPORewardsPublisher(RollbackAwarePublisher<Message>);
 
 impl SPORewardsPublisher {
     /// Construct with context and topic to publish on
     pub fn new(context: Arc<Context<Message>>, topic: String) -> Self {
-        Self {
-            context,
-            topic,
-            last_activity_at: None,
-        }
+        Self(RollbackAwarePublisher::new(context, topic))
     }
 
     /// Publish the SPO rewards
@@ -31,31 +19,19 @@ impl SPORewardsPublisher {
         block: &BlockInfo,
         spo_rewards: Vec<(PoolId, SPORewards)>,
     ) -> anyhow::Result<()> {
-        self.last_activity_at = Some(block.slot);
-        self.context
-            .message_bus
-            .publish(
-                &self.topic,
-                Arc::new(Message::Cardano((
-                    block.clone(),
-                    CardanoMessage::SPORewards(SPORewardsMessage {
-                        epoch: block.epoch - 1, // End of previous epoch
-                        spos: spo_rewards.into_iter().collect(),
-                    }),
-                ))),
-            )
+        self.0
+            .publish(Arc::new(Message::Cardano((
+                block.clone(),
+                CardanoMessage::SPORewards(SPORewardsMessage {
+                    epoch: block.epoch - 1, // End of previous epoch
+                    spos: spo_rewards.into_iter().collect(),
+                }),
+            ))))
             .await
     }
 
     /// Publish a rollback message, if we have anything to roll back
     pub async fn publish_rollback(&mut self, message: Arc<Message>) -> anyhow::Result<()> {
-        let Message::Cardano((block_info, CardanoMessage::Rollback(_))) = message.as_ref() else {
-            return Ok(());
-        };
-        if self.last_activity_at.is_none_or(|slot| slot < block_info.slot) {
-            return Ok(());
-        }
-        self.last_activity_at = None;
-        self.context.message_bus.publish(&self.topic, message).await
+        self.0.publish(message).await
     }
 }