Remove unnecessary counters from persistent state (#2639)

The `StableBTreeMap` tracks the number of elements on its own, so there is no need for an external counter. Removing the extra counters also makes the stats resetting obsolete and hence allows to remove the pre_stats II version from the integration tests.
dfinity · Oct 2, 2024 · 34c1321 · 34c1321
1 parent ee07a98
commit 34c1321
Show file tree

Hide file tree

Showing 7 changed files with 7 additions and 174 deletions.
diff --git a/.github/workflows/canister-tests.yml b/.github/workflows/canister-tests.yml
@@ -383,7 +383,6 @@ jobs:
           # a relatively small price to pay to make sure PRs are always tested against the latest release.
           curl -sSL https://github.com/dfinity/internet-identity/releases/latest/download/internet_identity_test.wasm.gz -o internet_identity_previous.wasm.gz
           curl -sSL https://github.com/dfinity/internet-identity/releases/latest/download/archive.wasm.gz -o archive_previous.wasm.gz
-          curl -SL https://github.com/dfinity/internet-identity/releases/download/release-2024-04-26/internet_identity_test.wasm.gz -o internet_identity_pre_stats.wasm.gz
 
           # We are using --partition hash instead of count, because it makes sure that the tests partition is stable across runs
           # even if tests are added or removed. The tradeoff is that the balancing might be slightly worse, but we have enough

diff --git a/src/canister_tests/src/framework.rs b/src/canister_tests/src/framework.rs
@@ -69,20 +69,6 @@ lazy_static! {
         get_wasm_path("II_WASM_PREVIOUS".to_string(), &def_path).expect(&err)
     };
 
-        /** The gzipped Wasm module for the II release build 2024-04-26, before the event stats were introduced. */
-    pub static ref II_WASM_PRE_STATS: Vec<u8> = {
-        let def_path = path::PathBuf::from("..").join("..").join("internet_identity_pre_stats.wasm.gz");
-        let err = format!("
-        Could not find Internet Identity Wasm module for pre-stats release.
-
-        I will look for it at {:?}, and you can specify another path with the environment variable II_WASM_PREVIOUS (note that I run from {:?}).
-
-        In order to get the Wasm module, please run the following command:
-            curl -SL https://github.com/dfinity/internet-identity/releases/download/release-2024-04-26/internet_identity_test.wasm.gz -o internet_identity_pre_stats.wasm.gz
-        ", &def_path, &std::env::current_dir().map(|x| x.display().to_string()).unwrap_or_else(|_| "an unknown directory".to_string()));
-        get_wasm_path("II_WASM_PRE_STATS".to_string(), &def_path).expect(&err)
-    };
-
         /** The gzipped Wasm module for the _previous_ archive build, or latest release, which is used when testing
             * upgrades and downgrades */
     pub static ref ARCHIVE_WASM_PREVIOUS: Vec<u8> = {

diff --git a/src/internet_identity/src/state.rs b/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());
     });
 }
 

diff --git a/src/internet_identity/src/stats/event_stats.rs b/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>,
 {

diff --git a/src/internet_identity/src/stats/event_stats/tests.rs b/src/internet_identity/src/stats/event_stats/tests.rs
@@ -75,8 +75,6 @@ fn should_store_event_and_add_to_aggregations() {
             .unwrap(),
         event
     );
-
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -98,7 +96,6 @@ fn should_store_multiple_events_with_same_timestamp() {
         assert_eq!(key.time, TIMESTAMP);
         assert_eq!(value, event.clone());
     });
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -149,7 +146,6 @@ fn should_track_ii_domains() {
     for key in expected_aggregations.iter() {
         assert!(storage.event_aggregations.contains_key(key));
     }
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -191,7 +187,6 @@ fn should_track_multiple_frontends() {
     for key in expected_aggregations.iter() {
         assert!(storage.event_aggregations.contains_key(key));
     }
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -258,7 +253,6 @@ fn should_store_multiple_events_and_aggregate_expected_weight_count() {
             .unwrap(),
         SESS_DURATION_SEC * 2
     );
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -273,9 +267,7 @@ fn should_remove_daily_events_after_24h() {
     };
 
     update_events_internal(event.clone(), TIMESTAMP, &mut storage);
-    assert_event_count_consistent(&mut storage);
     update_events_internal(event.clone(), TIMESTAMP + DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(storage.event_data.len(), 2);
     assert_eq!(storage.event_aggregations.len(), 4);
@@ -341,9 +333,7 @@ fn should_prune_monthly_events_after_30d() {
     };
 
     update_events_internal(event.clone(), TIMESTAMP, &mut storage);
-    assert_event_count_consistent(&mut storage);
     update_events_internal(event.clone(), TIMESTAMP + DAY_NS * 30, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(storage.event_data.len(), 1);
     assert_eq!(storage.event_aggregations.len(), 4);
@@ -419,7 +409,6 @@ fn should_remove_at_most_100_events_24h() {
     }
 
     update_events_internal(event.clone(), TIMESTAMP + DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     // Of the 107 initial events, 100 should be removed, 1 was added to trigger the clean-up
     // --> 8 expected events
@@ -436,7 +425,6 @@ fn should_remove_at_most_100_events_24h() {
         }
     );
     update_events_internal(event.clone(), TIMESTAMP + 2 * DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
     // Clean-up again, after another 24h leaving only the event that triggered the clean-up
     // --> 1 expected events
     assert_eq!(storage.event_aggregations.get(&aggregation_key).unwrap(), 1);
@@ -473,14 +461,12 @@ fn should_prune_at_most_100_events_30d() {
     }
 
     update_events_internal(event.clone(), TIMESTAMP + 30 * DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     // Of the 107 initial events, 100 should be pruned, 1 was added to trigger the pruning
     // --> 8 expected events
     assert_eq!(storage.event_aggregations.get(&aggregation_key).unwrap(), 8);
     assert_eq!(storage.event_data.len(), 8);
     update_events_internal(event.clone(), TIMESTAMP + 60 * DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
     // Prune again, after another 30d leaving only the event that triggered the pruning
     // --> 1 expected events
     assert_eq!(storage.event_aggregations.get(&aggregation_key).unwrap(), 1);
@@ -498,7 +484,6 @@ fn should_account_for_dapps_changing_session_lifetime() {
         }),
     };
     update_events_internal(event1, TIMESTAMP, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(
         storage
@@ -521,7 +506,6 @@ fn should_account_for_dapps_changing_session_lifetime() {
         }),
     };
     update_events_internal(event2, TIMESTAMP + 1, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(
         storage
@@ -544,7 +528,6 @@ fn should_account_for_dapps_changing_session_lifetime() {
         }),
     };
     update_events_internal(event3, TIMESTAMP + 1 + DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(
         storage
@@ -572,7 +555,6 @@ fn should_remove_aggregations_without_events_when_pruning() {
     };
 
     update_events_internal(event, TIMESTAMP, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     assert_eq!(storage.event_aggregations.len(), 4);
     let expected_aggregations: [AggregationKey; 4] = [
@@ -616,7 +598,6 @@ fn should_remove_aggregations_without_events_when_pruning() {
     // after this update, the previous aggregations should be removed as their weight is 0
     // (this means that there is no event being counted for them)
     update_events_internal(event2, TIMESTAMP + 30 * DAY_NS, &mut storage);
-    assert_event_count_consistent(&mut storage);
 
     let expected_aggregations_2: [AggregationKey; 4] = [
         AggregationKey::new(
@@ -706,7 +687,6 @@ fn should_ignore_prune_events_for_prepare_delegation_aggregations() {
     for key in pd_count_aggregations.iter() {
         assert_eq!(storage.event_aggregations.get(key).unwrap(), 1);
     }
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -724,7 +704,6 @@ fn should_add_prune_events_to_prune_aggregations() {
     for key in sess_sec_aggregations.iter() {
         assert_eq!(storage.event_aggregations.get(key).unwrap(), 3);
     }
-    assert_event_count_consistent(&mut storage);
 }
 
 #[test]
@@ -801,18 +780,6 @@ fn should_wrap_event_key_counter_correctly() {
     assert!(next_key > key);
 }
 
-/// Make sure the cached count values are consistent with the actual data
-fn assert_event_count_consistent(storage: &mut Storage<Rc<RefCell<Vec<u8>>>>) {
-    assert_eq!(
-        storage.event_data.iter().count(),
-        persistent_state(|s| s.event_data_count) as usize
-    );
-    assert_eq!(
-        storage.event_aggregations.iter().count(),
-        persistent_state(|s| s.event_aggregations_count) as usize
-    );
-}
-
 fn to_ns(secs: u64) -> u64 {
     Duration::from_secs(secs).as_nanos() as u64
 }

diff --git a/src/internet_identity/src/storage/storable_persistent_state.rs b/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,
@@ -154,8 +145,6 @@ mod tests {
                 max_unsolved_captchas: 500,
                 captcha_trigger: CaptchaTrigger::Static(StaticCaptchaTrigger::CaptchaEnabled),
             },
-            event_data_count: 0,
-            event_aggregations_count: 0,
             event_stats_24h_start: None,
         };
         assert_eq!(PersistentState::default(), expected_defaults);