fix/CRV3 Improvements

pallets/drand/src/lib.rs

@@ -58,6 +58,7 @@ use sp_runtime::{
 };
 
 pub mod bls12_381;
+pub mod migrations;
 pub mod types;
 pub mod utils;
 pub mod verifier;
@@ -91,6 +92,8 @@ pub const QUICKNET_CHAIN_HASH: &str =
 const CHAIN_HASH: &str = QUICKNET_CHAIN_HASH;
 
 pub const MAX_PULSES_TO_FETCH: u64 = 50;
+pub const MAX_KEPT_PULSES: u64 = 216_000; // 1 week
+pub const MAX_REMOVED_PULSES: u64 = 100;
 
 /// Defines application identifier for crypto keys of this module.
 ///
@@ -212,13 +215,25 @@ pub mod pallet {
         }
     }
 
+    /// Define a maximum length for the migration key
+    type MigrationKeyMaxLen = ConstU32<128>;
+
+    /// Storage for migration run status
+    #[pallet::storage]
+    pub type HasMigrationRun<T: Config> =
+        StorageMap<_, Identity, BoundedVec<u8, MigrationKeyMaxLen>, bool, ValueQuery>;
+
     /// map round number to pulse
     #[pallet::storage]
     pub type Pulses<T: Config> = StorageMap<_, Blake2_128Concat, RoundNumber, Pulse, OptionQuery>;
 
     #[pallet::storage]
     pub(super) type LastStoredRound<T: Config> = StorageValue<_, RoundNumber, ValueQuery>;
 
+    /// oldest stored round
+    #[pallet::storage]
+    pub(super) type OldestStoredRound<T: Config> = StorageValue<_, RoundNumber, ValueQuery>;
+
     /// Defines the block when next unsigned transaction will be accepted.
     ///
     /// To prevent spam of unsigned (and unpaid!) transactions on the network,
@@ -261,6 +276,14 @@ pub mod pallet {
                 log::debug!("Drand: Failed to fetch pulse from drand. {:?}", e);
             }
         }
+        fn on_runtime_upgrade() -> frame_support::weights::Weight {
+            /* let weight = */
+            frame_support::weights::Weight::from_parts(0, 0) /*;*/
+
+            //weight = weight.saturating_add(migrations::migrate_prune_old_pulses::<T>());
+
+            //weight
+        }
     }
 
     #[pallet::validate_unsigned]
@@ -308,8 +331,8 @@ pub mod pallet {
         /// Verify and write a pulse from the beacon into the runtime
         #[pallet::call_index(0)]
         #[pallet::weight(Weight::from_parts(5_708_000_000, 0)
-        .saturating_add(T::DbWeight::get().reads(2_u64))
-        .saturating_add(T::DbWeight::get().writes(3_u64)))]
+        .saturating_add(T::DbWeight::get().reads(3_u64))
+        .saturating_add(T::DbWeight::get().writes(4_u64)))]
         pub fn write_pulse(
             origin: OriginFor<T>,
             pulses_payload: PulsesPayload<T::Public, BlockNumberFor<T>>,
@@ -321,6 +344,10 @@ pub mod pallet {
             let mut last_stored_round = LastStoredRound::<T>::get();
             let mut new_rounds = Vec::new();
 
+            let oldest_stored_round = OldestStoredRound::<T>::get();
+            let is_first_storage = last_stored_round == 0 && oldest_stored_round == 0;
+            let mut first_new_round: Option<RoundNumber> = None;
+
             for pulse in &pulses_payload.pulses {
                 let is_verified = T::Verifier::verify(config.clone(), pulse.clone())
                     .map_err(|_| Error::<T>::PulseVerificationError)?;
@@ -339,12 +366,25 @@ pub mod pallet {
 
                     // Collect the new round
                     new_rounds.push(pulse.round);
+
+                    // Set the first new round if this is the initial storage
+                    if is_first_storage && first_new_round.is_none() {
+                        first_new_round = Some(pulse.round);
+                    }
                 }
             }
 
             // Update LastStoredRound storage
             LastStoredRound::<T>::put(last_stored_round);
 
+            // Set OldestStoredRound if this was the first storage
+            if let Some(first_round) = first_new_round {
+                OldestStoredRound::<T>::put(first_round);
+            }
+
+            // Prune old pulses
+            Self::prune_old_pulses(last_stored_round);
+
             // Update the next unsigned block number
             let current_block = frame_system::Pallet::<T>::block_number();
             <NextUnsignedAt<T>>::put(current_block);
@@ -628,6 +668,24 @@ impl<T: Config> Pallet<T> {
             .propagate(true)
             .build()
     }
+
+    fn prune_old_pulses(last_stored_round: RoundNumber) {
+        let mut oldest = OldestStoredRound::<T>::get();
+        if oldest == 0 {
+            return;
+        }
+
+        let mut removed: u64 = 0;
+        while last_stored_round.saturating_sub(oldest) + 1 > MAX_KEPT_PULSES
+            && removed < MAX_REMOVED_PULSES
+        {
+            Pulses::<T>::remove(oldest);
+            oldest = oldest.saturating_add(1);
+            removed = removed.saturating_add(1);
+        }
+
+        OldestStoredRound::<T>::put(oldest);
+    }
 }
 
 /// construct a message (e.g. signed by drand)

pallets/drand/src/migrations/migrate_prune_old_pulses.rs

@@ -0,0 +1,63 @@
+use crate::*;
+use frame_support::{traits::Get, weights::Weight};
+use log;
+
+pub fn migrate_prune_old_pulses<T: Config>() -> Weight {
+    let migration_name = BoundedVec::truncate_from(b"migrate_prune_old_pulses".to_vec());
+
+    // Initialize the weight with one read operation.
+    let mut weight = T::DbWeight::get().reads(1);
+
+    // Check if the migration has already run
+    if HasMigrationRun::<T>::get(&migration_name) {
+        log::info!(
+            "Migration '{:?}' has already run. Skipping.",
+            String::from_utf8_lossy(&migration_name)
+        );
+        return weight;
+    }
+    log::info!(
+        "Running migration '{}'",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // Collect all round numbers
+    let mut rounds: Vec<RoundNumber> = Pulses::<T>::iter_keys().collect();
+    weight = weight.saturating_add(T::DbWeight::get().reads(rounds.len() as u64));
+
+    if rounds.is_empty() {
+        OldestStoredRound::<T>::put(0u64);
+        LastStoredRound::<T>::put(0u64);
+        weight = weight.saturating_add(T::DbWeight::get().writes(2));
+    } else {
+        rounds.sort();
+        let num_pulses = rounds.len() as u64;
+
+        let mut new_oldest = rounds[0];
+        if num_pulses > MAX_KEPT_PULSES {
+            let num_to_delete = num_pulses - MAX_KEPT_PULSES;
+            new_oldest = rounds[num_to_delete as usize];
+
+            for &round in &rounds[0..num_to_delete as usize] {
+                Pulses::<T>::remove(round);
+                weight = weight.saturating_add(T::DbWeight::get().writes(1));
+            }
+        }
+
+        OldestStoredRound::<T>::put(new_oldest);
+        LastStoredRound::<T>::put(*rounds.last().unwrap());
+        weight = weight.saturating_add(T::DbWeight::get().writes(2));
+    }
+
+    // Mark the migration as completed
+    HasMigrationRun::<T>::insert(&migration_name, true);
+    weight = weight.saturating_add(T::DbWeight::get().writes(1));
+
+    log::info!(
+        "Migration '{:?}' completed.",
+        String::from_utf8_lossy(&migration_name)
+    );
+
+    // Return the migration weight.
+    weight
+}

pallets/drand/src/migrations/mod.rs

@@ -0,0 +1,2 @@
+pub mod migrate_prune_old_pulses;
+pub use migrate_prune_old_pulses::*;

pallets/drand/src/tests.rs

@@ -16,14 +16,17 @@
 
 use crate::{
     BeaconConfig, BeaconConfigurationPayload, BeaconInfoResponse, Call, DrandResponseBody,
-    ENDPOINTS, Error, Pulse, Pulses, PulsesPayload, QUICKNET_CHAIN_HASH, mock::*,
+    ENDPOINTS, Error, HasMigrationRun, LastStoredRound, MAX_KEPT_PULSES, OldestStoredRound, Pulse,
+    Pulses, PulsesPayload, QUICKNET_CHAIN_HASH, migrations::migrate_prune_old_pulses, mock::*,
 };
 use codec::Encode;
 use frame_support::{
-    assert_noop, assert_ok,
+    BoundedVec, assert_noop, assert_ok,
     pallet_prelude::{InvalidTransaction, TransactionSource},
+    weights::RuntimeDbWeight,
 };
 use frame_system::RawOrigin;
+use sp_core::Get;
 use sp_runtime::{
     offchain::{
         OffchainWorkerExt,
@@ -549,3 +552,158 @@ fn test_invalid_json_then_success() {
         assert_eq!(actual, expected_pulse);
     });
 }
+
+#[test]
+fn test_pulses_are_correctly_pruned() {
+    new_test_ext().execute_with(|| {
+        let pulse = Pulse::default();
+        let last_round: u64 = MAX_KEPT_PULSES + 2;
+        let oldest_round: u64 = 1;
+        let prune_count: u64 = 2;
+        let new_oldest: u64 = oldest_round + prune_count;
+        let middle_round: u64 = MAX_KEPT_PULSES / 2;
+
+        // Set storage bounds
+        OldestStoredRound::<Test>::put(oldest_round);
+        LastStoredRound::<Test>::put(last_round);
+
+        // Insert pulses at boundaries
+        // These should be pruned
+        Pulses::<Test>::insert(1, pulse.clone());
+        Pulses::<Test>::insert(2, pulse.clone());
+
+        // This should remain (new oldest)
+        Pulses::<Test>::insert(new_oldest, pulse.clone());
+
+        // Middle and last should remain
+        Pulses::<Test>::insert(middle_round, pulse.clone());
+        Pulses::<Test>::insert(last_round, pulse.clone());
+
+        // Trigger prune
+        Drand::prune_old_pulses(last_round);
+
+        // Assert new oldest
+        assert_eq!(OldestStoredRound::<Test>::get(), new_oldest);
+
+        // Assert pruned correctly
+        assert!(!Pulses::<Test>::contains_key(1), "Round 1 should be pruned");
+        assert!(!Pulses::<Test>::contains_key(2), "Round 2 should be pruned");
+
+        // Assert not pruned incorrectly
+        assert!(
+            Pulses::<Test>::contains_key(new_oldest),
+            "New oldest round should remain"
+        );
+        assert!(
+            Pulses::<Test>::contains_key(middle_round),
+            "Middle round should remain"
+        );
+        assert!(
+            Pulses::<Test>::contains_key(last_round),
+            "Last round should remain"
+        );
+    });
+}
+
+#[test]
+fn test_migrate_prune_old_pulses() {
+    new_test_ext().execute_with(|| {
+        let migration_name = BoundedVec::truncate_from(b"migrate_prune_old_pulses".to_vec());
+        let pulse = Pulse::default();
+
+        assert_eq!(Pulses::<Test>::iter().count(), 0);
+        assert!(!HasMigrationRun::<Test>::get(&migration_name));
+        assert_eq!(OldestStoredRound::<Test>::get(), 0);
+        assert_eq!(LastStoredRound::<Test>::get(), 0);
+
+        // Test with more pulses than MAX_KEPT_PULSES
+        let excess: u64 = 9;
+        let total: u64 = MAX_KEPT_PULSES + excess;
+        for i in 1..=total {
+            Pulses::<Test>::insert(i, pulse.clone());
+        }
+
+        let weight_large = migrate_prune_old_pulses::<Test>();
+
+        let expected_oldest = excess + 1;
+        assert_eq!(OldestStoredRound::<Test>::get(), expected_oldest);
+        assert_eq!(LastStoredRound::<Test>::get(), total);
+
+        for i in 1..=excess {
+            assert!(!Pulses::<Test>::contains_key(i));
+        }
+        for i in expected_oldest..=total {
+            assert!(Pulses::<Test>::contains_key(i));
+        }
+
+        let db_weight: RuntimeDbWeight = <Test as frame_system::Config>::DbWeight::get();
+        let num_pulses = total;
+        let num_to_delete = num_pulses - MAX_KEPT_PULSES;
+        let expected_weight = db_weight.reads(1 + num_pulses) + db_weight.writes(num_to_delete + 3);
+        assert_eq!(weight_large, expected_weight);
+    });
+}
+
+#[test]
+fn test_prune_maximum_of_100_pulses_per_call() {
+    new_test_ext().execute_with(|| {
+        // ------------------------------------------------------------
+        // 1. Arrange – create a storage layout that exceeds MAX_KEPT_PULSES
+        // ------------------------------------------------------------
+        const EXTRA: u64 = 250;
+        let oldest_round: u64 = 1;
+        let last_round: u64 = oldest_round + MAX_KEPT_PULSES + EXTRA;
+
+        OldestStoredRound::<Test>::put(oldest_round);
+        LastStoredRound::<Test>::put(last_round);
+        let pulse = Pulse::default();
+
+        // Insert the first 150 rounds so we can check they disappear / stay
+        for r in oldest_round..=oldest_round + 150 {
+            Pulses::<Test>::insert(r, pulse.clone());
+        }
+        let mid_round = oldest_round + 150;
+        Pulses::<Test>::insert(last_round, pulse.clone());
+
+        // ------------------------------------------------------------
+        // 2. Act – run the pruning function once
+        // ------------------------------------------------------------
+        Drand::prune_old_pulses(last_round);
+
+        // ------------------------------------------------------------
+        // 3. Assert – only the *first* 100 pulses were removed
+        // ------------------------------------------------------------
+        let expected_new_oldest = oldest_round + 100; // 101
+
+        // ‣ Storage bound updated correctly
+        assert_eq!(
+            OldestStoredRound::<Test>::get(),
+            expected_new_oldest,
+            "OldestStoredRound should advance by exactly 100"
+        );
+
+        // ‣ Rounds 1‑100 are gone
+        for r in oldest_round..expected_new_oldest {
+            assert!(
+                !Pulses::<Test>::contains_key(r),
+                "Round {} should have been pruned",
+                r
+            );
+        }
+
+        // ‣ Round 101 (new oldest) and later rounds remain
+        assert!(
+            Pulses::<Test>::contains_key(expected_new_oldest),
+            "Round {} should remain after pruning",
+            expected_new_oldest
+        );
+        assert!(
+            Pulses::<Test>::contains_key(mid_round),
+            "Mid‑range round should remain after pruning"
+        );
+        assert!(
+            Pulses::<Test>::contains_key(last_round),
+            "LastStoredRound should remain after pruning"
+        );
+    });
+}