Reset event data on empty persistent state counts (#2496)

* Reset event data on empty persistent state

This is necessary because the persistent state will reinitialize
the event and aggregation count to 0, if rolled back across a version
that did not have the event data feature yet.

* Add integration test

* Switch to saturating arithmetics for event data count

.github/workflows/canister-tests.yml

@@ -390,6 +390,7 @@ 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

src/canister_tests/src/framework.rs

@@ -70,6 +70,20 @@ 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> = {

src/internet_identity/src/state.rs

@@ -264,7 +264,17 @@ pub fn save_persistent_state() {
 
 pub fn load_persistent_state() {
     STATE.with(|s| {
-        storage_borrow(|storage| *s.persistent_state.borrow_mut() = storage.read_persistent_state())
+        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;
     });
 }
 

src/internet_identity/src/stats/event_stats.rs

@@ -208,7 +208,7 @@ fn update_events_internal<M: Memory>(event: EventData, now: Timestamp, s: &mut S
         ic_cdk::println!("WARN: Event already exists for key {:?}", current_key);
     }
     state::persistent_state_mut(|s| {
-        s.event_data_count += 1;
+        s.event_data_count = s.event_data_count.saturating_add(1);
     });
 
     let mut aggregations_db_wrapper = CountingAggregationsWrapper(&mut s.event_aggregations);
@@ -344,7 +344,9 @@ fn prune_events<M: Memory>(
         db.remove(&entry.0);
     }
     state::persistent_state_mut(|s| {
-        s.event_data_count -= pruned_events.len() as u64;
+        s.event_data_count = s
+            .event_data_count
+            .saturating_sub(pruned_events.len() as u64);
     });
     pruned_events
 }
@@ -358,7 +360,7 @@ impl<'a, M: Memory> CountingAggregationsWrapper<'a, M> {
         if prev_value.is_none() {
             // Increase count because we added a new value
             state::persistent_state_mut(|s| {
-                s.event_aggregations_count += 1;
+                s.event_aggregations_count = s.event_aggregations_count.saturating_add(1);
             })
         }
     }
@@ -372,7 +374,7 @@ impl<'a, M: Memory> CountingAggregationsWrapper<'a, M> {
         if prev_value.is_some() {
             // Decrease count because we removed a value
             state::persistent_state_mut(|s| {
-                s.event_aggregations_count -= 1;
+                s.event_aggregations_count = s.event_aggregations_count.saturating_sub(1);
             })
         }
     }

src/internet_identity/tests/integration/aggregation_stats.rs

@@ -4,7 +4,7 @@ use crate::v2_api::authn_method_test_helpers::{
 use canister_tests::api::internet_identity as api;
 use canister_tests::framework::{
     assert_metric, env, get_metrics, install_ii_canister, restore_compressed_stable_memory,
-    upgrade_ii_canister, EMPTY_WASM, II_WASM,
+    upgrade_ii_canister, EMPTY_WASM, II_WASM, II_WASM_PRE_STATS,
 };
 use ic_cdk::api::management_canister::main::CanisterId;
 use ic_test_state_machine_client::{CallError, StateMachine};
@@ -264,6 +264,56 @@ fn should_keep_aggregations_across_upgrades() -> Result<(), CallError> {
     Ok(())
 }
 
+#[test]
+fn should_reset_stats_on_rollback_and_upgrade() -> Result<(), CallError> {
+    const II_ORIGIN: &str = "ic0.app";
+
+    let env = env();
+    let canister_id = install_ii_canister(&env, II_WASM.clone());
+    let identity_nr = create_identity(&env, canister_id, II_ORIGIN);
+
+    delegation_for_origin(&env, canister_id, identity_nr, "https://some-dapp.com")?;
+    delegation_for_origin(&env, canister_id, identity_nr, "https://some-dapp.com")?;
+
+    let aggregations = api::stats(&env, canister_id)?.event_aggregations;
+    assert_expected_aggregation(
+        &aggregations,
+        &aggregation_key(PD_COUNT, "24h", II_ORIGIN),
+        vec![("https://some-dapp.com".to_string(), 2u64)],
+    );
+    assert_expected_aggregation(
+        &aggregations,
+        &aggregation_key(PD_COUNT, "30d", II_ORIGIN),
+        vec![("https://some-dapp.com".to_string(), 2u64)],
+    );
+    assert_expected_aggregation(
+        &aggregations,
+        &aggregation_key(PD_SESS_SEC, "24h", II_ORIGIN),
+        vec![("https://some-dapp.com".to_string(), 2 * SESSION_LENGTH)],
+    );
+    assert_expected_aggregation(
+        &aggregations,
+        &aggregation_key(PD_SESS_SEC, "30d", II_ORIGIN),
+        vec![("https://some-dapp.com".to_string(), 2 * SESSION_LENGTH)],
+    );
+    let metrics = get_metrics(&env, canister_id);
+    assert_metric(&metrics, "internet_identity_event_data_count", 2f64);
+    assert_metric(&metrics, "internet_identity_event_aggregations_count", 4f64);
+
+    // roll back to a version that does not know about event stats
+    upgrade_ii_canister(&env, canister_id, II_WASM_PRE_STATS.clone());
+    // upgrade to the latest version again --> stats should be reset
+    upgrade_ii_canister(&env, canister_id, II_WASM.clone());
+
+    let aggregations = api::stats(&env, canister_id)?.event_aggregations;
+    assert_eq!(aggregations.len(), 0);
+
+    let metrics = get_metrics(&env, canister_id);
+    assert_metric(&metrics, "internet_identity_event_data_count", 0f64);
+    assert_metric(&metrics, "internet_identity_event_aggregations_count", 0f64);
+    Ok(())
+}
+
 #[test]
 fn crash_mem_backup() -> Result<(), CallError> {
     let env = env();