Do not trap on pruning non-existing aggregations or zero weighted events

This PR fixes a bug introduced in #2449, which subsequently led to a
rollback in proposal 130068.

This refactors the code to accommodate violated expectations and proceed
without trapping.
Zero weighted event are now correctly ignored by the pruning logic.

src/internet_identity/src/stats/event_stats.rs

@@ -389,6 +389,8 @@ fn update_aggregation<M: Memory, F>(
     let mut pruned_weight_by_key = pruned_events
         .iter()
         .filter_map(&aggregation_filter_map)
+        // zero weighted events are ignored for pruning
+        .filter(|weighted_aggregation| weighted_aggregation.weight > 0)
         .fold(HashMap::<_, u64>::new(), |mut map, weighted_aggregation| {
             let value =
                 map.get(&weighted_aggregation.key).unwrap_or(&0) + weighted_aggregation.weight;
@@ -420,21 +422,18 @@ fn update_aggregation<M: Memory, F>(
     pruned_weight_by_key
         .into_iter()
         .for_each(|(key, pruned_weight)| {
-            let current_weight = db.get(&key).unwrap_or_else(|| {
-                panic!("aggregation key \"{:?}\" not found in DB when pruning", key)
-            });
-            assert!(
-                current_weight >= pruned_weight,
-                "pruned weight {} exceeds current weight {} for key \"{:?}\"",
-                pruned_weight,
-                current_weight,
-                key
-            );
-            let new_weight = current_weight - pruned_weight;
-            if new_weight == 0 {
-                db.remove(&key);
+            // Adjust aggregation weight by subtracting the pruned weight
+            if let Some(current_weight) = db.get(&key) {
+                let new_weight = current_weight.saturating_sub(pruned_weight);
+                if new_weight == 0 {
+                    db.remove(&key);
+                } else {
+                    db.insert(key, new_weight);
+                }
             } else {
-                db.insert(key, new_weight);
+                // This should not happen, but if it does, it's a bug.
+                // Print a message and so that we can investigate once canister logs are available.
+                ic_cdk::println!("WARN: Aggregation key {:?} not found in db", key);
             }
         });
 }

src/internet_identity/src/stats/event_stats/tests.rs

@@ -637,6 +637,61 @@ fn should_add_prune_events_to_prune_aggregations() {
     assert_event_count_consistent(&mut storage);
 }
 
+#[test]
+fn should_prune_zero_weighted_events() {
+    let mut storage = test_storage();
+    let ii_domain = Some(IIDomain::Ic0App);
+    let event = EventData {
+        event: Event::PrepareDelegation(PrepareDelegationEvent {
+            ii_domain: ii_domain.clone(),
+            frontend: EXAMPLE_URL.to_string(),
+            session_duration_ns: to_ns(0),
+        }),
+    };
+
+    // no events initially
+    assert_eq!(storage.event_data.len(), 0);
+
+    // add multiple zero weighted event
+    update_events_internal(event.clone(), TIMESTAMP, &mut storage);
+    update_events_internal(event.clone(), TIMESTAMP, &mut storage);
+    update_events_internal(event.clone(), TIMESTAMP, &mut storage);
+
+    assert_eq!(storage.event_data.len(), 3);
+
+    // there should not be any aggregation for session seconds as the weight is 0
+    let sess_sec_aggregations: [AggregationKey; 2] = [
+        AggregationKey::new(
+            PrepareDelegationSessionSeconds,
+            Day,
+            ii_domain.clone(),
+            EXAMPLE_URL.to_string(),
+        ),
+        AggregationKey::new(
+            PrepareDelegationSessionSeconds,
+            Month,
+            ii_domain.clone(),
+            EXAMPLE_URL.to_string(),
+        ),
+    ];
+    for key in sess_sec_aggregations.iter() {
+        assert!(storage.event_aggregations.get(key).is_none());
+    }
+
+
+    let event2 = EventData {
+        event: Event::PrepareDelegation(PrepareDelegationEvent {
+            ii_domain: Some(IIDomain::Ic0App),
+            frontend: EXAMPLE_URL_2.to_string(),
+            session_duration_ns: to_ns(SESS_DURATION_SEC),
+        }),
+    };
+    update_events_internal(event2.clone(), TIMESTAMP + DAY_NS * 30, &mut storage);
+
+    // zero weighted events should have been pruned
+    assert_eq!(storage.event_data.len(), 1);
+}
+
 /// 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!(