Remove unnecessary counters from persistent state

The `StableBTreeMap` tracks the number of elements on its own, so there
is no need for an external counter.

src/internet_identity/src/state.rs

@@ -103,14 +103,6 @@ pub struct PersistentState {
     pub active_authn_method_stats: ActivityStats<AuthnMethodCounter>,
     // Configuration of the captcha challenge during registration flow
     pub captcha_config: CaptchaConfig,
-    // Count of entries in the event_data BTreeMap
-    // event_data is expected to have a lot of entries, thus counting by iterating over it is not
-    // an option.
-    pub event_data_count: u64,
-    // Count of entries in the event_aggregations BTreeMap
-    // event_aggregations is expected to have a lot of entries, thus counting by iterating over it is not
-    // an option.
-    pub event_aggregations_count: u64,
     // Key into the event_data BTreeMap where the 24h tracking window starts.
     // This key is used to remove old entries from the 24h event aggregations.
     // If it is `none`, then the 24h window starts from the newest entry in the event_data
@@ -129,8 +121,6 @@ impl Default for PersistentState {
             domain_active_anchor_stats: ActivityStats::new(time),
             active_authn_method_stats: ActivityStats::new(time),
             captcha_config: DEFAULT_CAPTCHA_CONFIG,
-            event_data_count: 0,
-            event_aggregations_count: 0,
             event_stats_24h_start: None,
         }
     }
@@ -257,17 +247,8 @@ pub fn save_persistent_state() {
 
 pub fn load_persistent_state() {
     STATE.with(|s| {
-        let loaded_state = storage_borrow(|storage| storage.read_persistent_state());
-
-        // Reset the event_data and event_aggregations if they are reported as empty from the persistent state
-        // Necessary to handle rollback across the versions where the event_data and event_aggregations were introduced
-        if loaded_state.event_aggregations_count == 0 {
-            storage_borrow_mut(|storage| storage.event_aggregations.clear_new());
-        }
-        if loaded_state.event_data_count == 0 {
-            storage_borrow_mut(|storage| storage.event_data.clear_new());
-        }
-        *s.persistent_state.borrow_mut() = loaded_state;
+        *s.persistent_state.borrow_mut() =
+            storage_borrow(|storage| storage.read_persistent_state());
     });
 }
 

src/internet_identity/src/stats/event_stats.rs

@@ -207,11 +207,7 @@ fn update_events_internal<M: Memory>(event: EventData, now: Timestamp, s: &mut S
     if option.is_some() {
         ic_cdk::println!("WARN: Event already exists for key {:?}", current_key);
     }
-    state::persistent_state_mut(|s| {
-        s.event_data_count = s.event_data_count.saturating_add(1);
-    });
 
-    let mut aggregations_db_wrapper = CountingAggregationsWrapper(&mut s.event_aggregations);
     // Collect events that should no longer be reflected in the 24h aggregations. Does _not_ delete
     // the underlying events.
     let removed_24h = events_to_remove_from_24h_aggregations(now, &current_key, &s.event_data);
@@ -221,7 +217,7 @@ fn update_events_internal<M: Memory>(event: EventData, now: Timestamp, s: &mut S
         &event,
         &removed_24h,
         AggregationWindow::Day,
-        &mut aggregations_db_wrapper,
+        &mut s.event_aggregations,
     );
 
     // This pruning _deletes_ the data older than 30 days. Do this after the 24h aggregation
@@ -234,7 +230,7 @@ fn update_events_internal<M: Memory>(event: EventData, now: Timestamp, s: &mut S
         &event,
         &pruned_30d,
         AggregationWindow::Month,
-        &mut aggregations_db_wrapper,
+        &mut s.event_aggregations,
     );
 }
 
@@ -244,7 +240,7 @@ fn update_aggregations<M: Memory>(
     event_data: &EventData,
     pruned_events: &[(EventKey, EventData)],
     window: AggregationWindow,
-    aggregations_db: &mut CountingAggregationsWrapper<M>,
+    aggregations_db: &mut StableBTreeMap<AggregationKey, u64, M>,
 ) {
     AGGREGATIONS.iter().for_each(|aggregation| {
         update_aggregation(
@@ -343,49 +339,14 @@ fn prune_events<M: Memory>(
         let entry: &(EventKey, EventData) = entry;
         db.remove(&entry.0);
     }
-    state::persistent_state_mut(|s| {
-        s.event_data_count = s
-            .event_data_count
-            .saturating_sub(pruned_events.len() as u64);
-    });
     pruned_events
 }
-
-/// Helper struct for updating the count of event aggregations.
-struct CountingAggregationsWrapper<'a, M: Memory>(&'a mut StableBTreeMap<AggregationKey, u64, M>);
-
-impl<'a, M: Memory> CountingAggregationsWrapper<'a, M> {
-    fn insert(&mut self, key: AggregationKey, value: u64) {
-        let prev_value = self.0.insert(key, value);
-        if prev_value.is_none() {
-            // Increase count because we added a new value
-            state::persistent_state_mut(|s| {
-                s.event_aggregations_count = s.event_aggregations_count.saturating_add(1);
-            })
-        }
-    }
-
-    fn get(&self, key: &AggregationKey) -> Option<u64> {
-        self.0.get(key)
-    }
-
-    fn remove(&mut self, key: &AggregationKey) {
-        let prev_value = self.0.remove(key);
-        if prev_value.is_some() {
-            // Decrease count because we removed a value
-            state::persistent_state_mut(|s| {
-                s.event_aggregations_count = s.event_aggregations_count.saturating_sub(1);
-            })
-        }
-    }
-}
-
 fn update_aggregation<M: Memory, F>(
     aggregation_filter_map: F,
     now: EventKey,
     new_event: EventData,
     pruned_events: &[(EventKey, EventData)],
-    db: &mut CountingAggregationsWrapper<M>,
+    db: &mut StableBTreeMap<AggregationKey, u64, M>,
 ) where
     F: Fn(&(EventKey, EventData)) -> Option<AggregationEvent>,
 {

src/internet_identity/src/storage/storable_persistent_state.rs

@@ -30,10 +30,7 @@ pub struct StorablePersistentState {
     max_inflight_captchas: u64,
 
     // opt fields because of backwards compatibility
-    event_data_count: Option<u64>,
-    event_aggregations_count: Option<u64>,
     event_stats_24h_start: Option<EventKey>,
-
     captcha_config: Option<CaptchaConfig>,
 }
 
@@ -70,8 +67,6 @@ impl From<PersistentState> for StorablePersistentState {
             max_num_latest_delegation_origins: 0,
             // unused, kept for stable memory compatibility
             max_inflight_captchas: 0,
-            event_data_count: Some(s.event_data_count),
-            event_aggregations_count: Some(s.event_aggregations_count),
             event_stats_24h_start: s.event_stats_24h_start,
             captcha_config: Some(s.captcha_config),
         }
@@ -88,8 +83,6 @@ impl From<StorablePersistentState> for PersistentState {
             domain_active_anchor_stats: s.domain_active_anchor_stats,
             active_authn_method_stats: s.active_authn_method_stats,
             captcha_config: s.captcha_config.unwrap_or(DEFAULT_CAPTCHA_CONFIG),
-            event_data_count: s.event_data_count.unwrap_or_default(),
-            event_aggregations_count: s.event_aggregations_count.unwrap_or_default(),
             event_stats_24h_start: s.event_stats_24h_start,
         }
     }
@@ -129,8 +122,6 @@ mod tests {
             latest_delegation_origins: HashMap::new(),
             max_num_latest_delegation_origins: 0,
             max_inflight_captchas: 0,
-            event_data_count: Some(0),
-            event_aggregations_count: Some(0),
             event_stats_24h_start: None,
             captcha_config: Some(CaptchaConfig {
                 max_unsolved_captchas: 500,