Simplify II init

A bug was introduced in #2460 that did not correctly restart timers
after an upgrade.
This PR simplifies the initialization to share more code and adds
a test to catch this bug in the future.

src/internet_identity/src/main.rs

@@ -393,13 +393,19 @@ fn inject_prune_event(timestamp: Timestamp) {
 #[init]
 #[candid_method(init)]
 fn init(maybe_arg: Option<InternetIdentityInit>) {
-    init_assets();
     state::init_new();
+    initialize(maybe_arg);
+}
 
-    apply_install_arg(maybe_arg);
+#[post_upgrade]
+fn post_upgrade(maybe_arg: Option<InternetIdentityInit>) {
+    state::init_from_stable_memory();
+    initialize(maybe_arg);
+}
 
-    // make sure the fully initialized storage configuration is written to stable memory
-    state::storage_borrow_mut(|storage| storage.flush());
+fn initialize(maybe_arg: Option<InternetIdentityInit>) {
+    init_assets();
+    apply_install_arg(maybe_arg);
     update_root_hash();
 
     // Immediately prune events if necessary and also set a timer to do so every 5 minutes.
@@ -412,19 +418,6 @@ fn init(maybe_arg: Option<InternetIdentityInit>) {
     });
 }
 
-#[post_upgrade]
-fn post_upgrade(maybe_arg: Option<InternetIdentityInit>) {
-    init_assets();
-    state::init_from_stable_memory();
-    // load the persistent state after initializing storage, otherwise the memory address to load it from cannot be calculated
-    state::load_persistent_state();
-
-    // We drop all the signatures on upgrade, users will
-    // re-request them if needed.
-    update_root_hash();
-    apply_install_arg(maybe_arg);
-}
-
 fn apply_install_arg(maybe_arg: Option<InternetIdentityInit>) {
     if let Some(arg) = maybe_arg {
         if let Some(range) = arg.assigned_user_number_range {

src/internet_identity/tests/integration/aggregation_stats.rs

@@ -122,19 +122,6 @@ fn should_report_at_most_100_entries() -> Result<(), CallError> {
     Ok(())
 }
 
-fn assert_expected_aggregation(
-    aggregations: &HashMap<String, Vec<(String, u64)>>,
-    key: &str,
-    data: Vec<(String, u64)>,
-) {
-    assert_eq!(
-        aggregations.get(key).unwrap(),
-        &data,
-        "Aggregation key \"{}\" does not match",
-        key
-    );
-}
-
 #[test]
 fn should_keep_aggregations_across_upgrades() -> Result<(), CallError> {
     const II_ORIGIN: &str = "ic0.app";
@@ -188,6 +175,49 @@ fn should_keep_aggregations_across_upgrades() -> Result<(), CallError> {
     Ok(())
 }
 
+#[test]
+fn should_prune_automatically_after_upgrade() -> Result<(), CallError> {
+    let env = env();
+    let canister_id = install_ii_canister(&env, II_WASM.clone());
+    let identity_nr = create_identity(&env, canister_id, "ic0.app");
+
+    delegation_for_origin(&env, canister_id, identity_nr, "https://some-dapp.com")?;
+
+    // upgrade then advance time to trigger pruning
+    upgrade_ii_canister(&env, canister_id, II_WASM.clone());
+    env.advance_time(Duration::from_secs(60 * 60 * 24 * 30)); // 30 days
+    env.tick(); // execute timers
+
+    let aggregations = api::stats(&env, canister_id)?.event_aggregations;
+
+    // set of keys only reflects pruning now
+    let mut expected_keys = vec![
+        aggregation_key(PRUNE_EVENT_COUNT, "24h", "other"),
+        aggregation_key(PRUNE_EVENT_COUNT, "30d", "other"),
+    ];
+    let mut keys = aggregations.into_keys().collect::<Vec<_>>();
+    // sort for stable comparison
+    expected_keys.sort();
+    keys.sort();
+    assert_eq!(keys, expected_keys);
+
+    Ok(())
+}
+
+
+fn assert_expected_aggregation(
+    aggregations: &HashMap<String, Vec<(String, u64)>>,
+    key: &str,
+    data: Vec<(String, u64)>,
+) {
+    assert_eq!(
+        aggregations.get(key).unwrap(),
+        &data,
+        "Aggregation key \"{}\" does not match",
+        key
+    );
+}
+
 fn create_identity(env: &StateMachine, canister_id: CanisterId, ii_origin: &str) -> IdentityNumber {
     let authn_method = AuthnMethodData {
         metadata: HashMap::from([(