Reduce memory usage in validator (#379)

Author: dknopik

anchor/client/src/lib.rs

@@ -416,14 +416,15 @@ impl Client {
         ));
         duties_tracker.clone().start(executor.clone());
 
-        let message_validator = Arc::new(Validator::new(
+        let message_validator = Validator::new(
             database.watch(),
             E::slots_per_epoch(),
             spec.epochs_per_sync_committee_period.as_u64(),
             E::sync_committee_size(),
             duties_tracker.clone(),
             slot_clock.clone(),
-        ));
+            &executor,
+        );
 
         let message_sender: Arc<dyn MessageSender> = if config.impostor.is_none() {
             Arc::new(NetworkMessageSender::new(

anchor/fuzz/fuzz_targets/setup.rs

@@ -116,14 +116,20 @@ pub fn setup_test_message_validator() -> Arc<Validator<ManualSlotClock, MockDuti
 
     let duties_provider = MockDutiesProvider {};
 
-    Arc::new(Validator::new(
+    let handle = tokio::runtime::Handle::current();
+    let (_signal, exit) = async_channel::bounded(1);
+    let (shutdown_tx, _) = futures::channel::mpsc::channel(1);
+    let executor = TaskExecutor::new(handle, exit, shutdown_tx, "test_executor".into());
+
+    Validator::new(
         db.watch(),
         32,
         256,
         512,
         duties_provider.into(),
         slot_clock.clone(),
-    ))
+        &executor,
+    )
 }
 
 // Sets up a real NetworkMessageReceiver for fuzzing
@@ -138,7 +144,7 @@ pub fn setup_test_message_receiver()
         max_workers: 2,
         queue_size: Default::default(),
     };
-    let processor_senders = processor::spawn(processor_config, executor);
+    let processor_senders = processor::spawn(processor_config, executor.clone());
 
     let slot_clock = ManualSlotClock::new(
         types::Slot::new(0),
@@ -161,14 +167,15 @@ pub fn setup_test_message_receiver()
 
     let duties_provider = MockDutiesProvider {};
 
-    let message_validator = Arc::new(Validator::new(
+    let message_validator = Validator::new(
         db.watch(),
         32,
         256,
         512,
         duties_provider.into(),
         slot_clock.clone(),
-    ));
+        &executor,
+    );
 
     let network_message_sender: Arc<dyn MessageSender> = Arc::new(
         NetworkMessageSender::new(

anchor/message_sender/src/network.rs

@@ -94,7 +94,7 @@ impl<S: SlotClock + 'static, D: DutiesProvider> MessageSender for Arc<NetworkMes
     }
 }
 
-impl<S: SlotClock, D: DutiesProvider> NetworkMessageSender<S, D> {
+impl<S: SlotClock + 'static, D: DutiesProvider> NetworkMessageSender<S, D> {
     pub fn new(
         processor: processor::Senders,
         network_tx: mpsc::Sender<(SubnetId, Vec<u8>)>,

anchor/message_validator/src/consensus_message.rs

@@ -221,12 +221,13 @@ pub(crate) fn validate_qbft_logic(
             }
 
             if consensus_message.round == signer_state.round {
-                // Rule: Peer must not send two proposals with different data
+                // Rule: Peer must not send two proposals with different data.
+                // We separately verify that the root in the message matches the data.
                 if !signed_ssv_message.full_data().is_empty()
                     && signer_state
-                        .proposal_data
+                        .proposal_hash
                         .as_ref()
-                        .is_some_and(|data| data != signed_ssv_message.full_data())
+                        .is_some_and(|hash| hash != consensus_message.root)
                 {
                     return Err(ValidationFailure::DifferentProposalData);
                 }

anchor/message_validator/src/duty_state.rs

@@ -104,6 +104,17 @@ impl DutyState {
 
         Ok(())
     }
+
+    /// Returns true if all operators within the map have a `max_slot` lower than `now -
+    /// stored_slot_count`. This indicates that there has been no relevant activity for this duty
+    /// recently and no relevant information is lost if this is dropped.
+    pub(crate) fn outdated(&self, current_slot: Slot) -> bool {
+        let earliest_relevant_slot =
+            current_slot.saturating_sub(Slot::from(self.stored_slot_count));
+        self.operators
+            .values()
+            .all(|operator_state| operator_state.max_slot < earliest_relevant_slot)
+    }
 }
 
 /// Tracks the state for a specific operator across multiple slots.
@@ -262,8 +273,8 @@ pub(crate) struct SignerState {
     pub(crate) round: u64,
     /// Records the count of each type of consensus message encountered.
     pub(crate) message_counts: MessageCounts,
-    /// Optionally holds proposal-related data if a proposal message was received.
-    pub(crate) proposal_data: Option<Vec<u8>>,
+    /// Holds the hash of the proposal data, if a proposal was received.
+    pub(crate) proposal_hash: Option<[u8; 32]>,
     /// A set of CommitteeIds indicating which committees have already been seen.
     seen_signers: HashSet<CommitteeId>,
 }
@@ -275,7 +286,7 @@ impl SignerState {
             slot,
             round,
             message_counts: MessageCounts::default(),
-            proposal_data: None,
+            proposal_hash: None,
             seen_signers: HashSet::new(),
         }
     }
@@ -289,14 +300,15 @@ impl SignerState {
 
     /// Updates the SignerState with a new consensus message.
     ///
-    /// - If the message is a proposal (and contains full data), it stores the proposal data.
+    /// - If the message is a proposal (and contains full data), it stores the hashed data.
     /// - If multiple operator IDs are present, it records the committee as seen.
     /// - Updates the message counts based on the message type.
     fn update(&mut self, signed_ssv_message: &SignedSSVMessage, consensus_message: &QbftMessage) {
         if !signed_ssv_message.full_data().is_empty()
             && consensus_message.qbft_message_type == QbftMessageType::Proposal
         {
-            self.proposal_data = Some(Vec::from(signed_ssv_message.full_data()));
+            // We verified that the proposal data matches the root.
+            self.proposal_hash = Some(*consensus_message.root);
         }
 
         if signed_ssv_message.operator_ids().len() > 1 {
@@ -316,18 +328,22 @@ mod tests {
     use ssv_types::{OperatorId, Slot, consensus::QbftMessageType, msgid::Role};
 
     use super::*;
-    use crate::tests::{QbftMessageBuilder, create_signed_consensus_message};
+    use crate::{
+        hash_data,
+        tests::{QbftMessageBuilder, create_signed_consensus_message},
+    };
 
     #[test]
     fn test_duty_state_update() {
         let mut duty_state = DutyState::new(10);
 
-        let qbft_message =
+        let mut qbft_message =
             QbftMessageBuilder::new(Role::Committee, QbftMessageType::Proposal).build();
 
         let operator_id = OperatorId(1);
 
         let full_data = vec![1, 2, 3];
+        *qbft_message.root = hash_data(&full_data);
         let signed_ssv_message = create_signed_consensus_message(
             qbft_message.clone(),
             vec![operator_id],
@@ -346,9 +362,9 @@ mod tests {
         if let Some(signer_state) = operator_state.get_signer_state(&slot) {
             // // Verify that the proposal data was correctly stored
             assert_eq!(
-                &signer_state.proposal_data,
-                &Some(full_data),
-                "Proposal data should match the signed message data"
+                signer_state.proposal_hash,
+                Some(hash_data(&full_data)),
+                "Proposal data should match the hashed full data"
             );
 
             // Verify message counts were updated

anchor/message_validator/src/lib.rs

@@ -29,7 +29,8 @@ use ssv_types::{
     partial_sig::PartialSignatureMessages,
 };
 use ssz::{Decode, DecodeError, Encode};
-use tokio::sync::watch::Receiver;
+use task_executor::TaskExecutor;
+use tokio::{sync::watch::Receiver, time::sleep};
 use tracing::{error, trace};
 use types::{Epoch, Slot};
 
@@ -39,6 +40,8 @@ use crate::{
     partial_signature::validate_partial_signature_message,
 };
 
+const VALIDATOR_CLEANER_NAME: &str = "validator_cleaner";
+
 pub(crate) const FIRST_ROUND: u64 = 1;
 
 #[derive(Debug)]
@@ -271,24 +274,29 @@ pub struct Validator<S: SlotClock, D: DutiesProvider> {
     slot_clock: S,
 }
 
-impl<S: SlotClock, D: DutiesProvider> Validator<S, D> {
+impl<S: SlotClock + 'static, D: DutiesProvider> Validator<S, D> {
     pub fn new(
         network_state_rx: Receiver<NetworkState>,
         slots_per_epoch: u64,
         epochs_per_sync_committee_period: u64,
         sync_committee_size: usize,
         duties_provider: Arc<D>,
         slot_clock: S,
-    ) -> Self {
-        Self {
+        task_executor: &TaskExecutor,
+    ) -> Arc<Self> {
+        let validator = Arc::new(Self {
             network_state_rx,
             duty_state_map: DashMap::new(),
             slots_per_epoch,
             epochs_per_sync_committee_period,
             sync_committee_size,
             duties_provider,
             slot_clock,
-        }
+        });
+
+        task_executor.spawn(Arc::clone(&validator).cleaner(), VALIDATOR_CLEANER_NAME);
+
+        validator
     }
 
     pub fn validate(&self, message_data: &[u8]) -> ValidationResult {
@@ -381,6 +389,40 @@ impl<S: SlotClock, D: DutiesProvider> Validator<S, D> {
                 DutyState::new(stored_slot_count as usize)
             })
     }
+
+    async fn cleaner(self: Arc<Self>) {
+        let slot_clock = self.slot_clock.clone();
+        let slots_per_epoch = self.slots_per_epoch;
+
+        // Use a weak reference to exit when the other `Arc` are dropped.
+        let weak_self = Arc::downgrade(&self);
+        loop {
+            // Try to get the time to the next slot.
+            let Some(until_next_epoch) = slot_clock.duration_to_next_epoch(slots_per_epoch) else {
+                sleep(slot_clock.slot_duration()).await;
+                continue;
+            };
+
+            // Wait until 5/6ths into the slot. Then, all proposal and attestation duties should be
+            // done, so we can lock the map without risking message delays for time-critical
+            // messages.
+            let sleep_for = until_next_epoch + slot_clock.slot_duration() * 5 / 6;
+            sleep(sleep_for).await;
+
+            let Some(validator) = weak_self.upgrade() else {
+                // No validator to clean anymore, exit.
+                break;
+            };
+            let Some(now) = slot_clock.now() else {
+                // Very weird, let's try again later.
+                continue;
+            };
+
+            validator
+                .duty_state_map
+                .retain(|_, duty_state| !duty_state.outdated(now));
+        }
+    }
 }
 
 fn validate_ssv_message(

anchor/message_validator/src/message_counts.rs

@@ -6,17 +6,17 @@ use ssv_types::{
 
 use crate::ValidationFailure;
 
-const MAX_MESSAGES_PER_ROUND: u64 = 1;
+const MAX_MESSAGES_PER_ROUND: u8 = 1;
 
 /// MessageCounts tracks different types of message counts per slot
 #[derive(Debug, Clone, Copy, Default)]
 pub(crate) struct MessageCounts {
-    pub(crate) pre_consensus: u64,
-    pub(crate) proposal: u64,
-    pub(crate) prepare: u64,
-    pub(crate) commit: u64,
-    pub(crate) round_change: u64,
-    pub(crate) post_consensus: u64,
+    pub(crate) pre_consensus: u8,
+    pub(crate) proposal: u8,
+    pub(crate) prepare: u8,
+    pub(crate) commit: u8,
+    pub(crate) round_change: u8,
+    pub(crate) post_consensus: u8,
 }
 
 impl MessageCounts {