diff --git a/src/internet_identity/src/stats/event_stats.rs b/src/internet_identity/src/stats/event_stats.rs index 961ebcde49..97fd6c4ab7 100644 --- a/src/internet_identity/src/stats/event_stats.rs +++ b/src/internet_identity/src/stats/event_stats.rs @@ -389,6 +389,10 @@ fn update_aggregation( let mut pruned_weight_by_key = pruned_events .iter() .filter_map(&aggregation_filter_map) + // zero weighted events are ignored when adjusting aggregations during pruning because they + // don't contribute any change to the aggregations hence pruning them has no effect on + // aggregations + .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 +424,18 @@ fn update_aggregation( 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 so that we can investigate once canister logs are available. + ic_cdk::println!("WARN: Aggregation key {:?} not found in db", key); } }); } diff --git a/src/internet_identity/src/stats/event_stats/tests.rs b/src/internet_identity/src/stats/event_stats/tests.rs index a56ce142d4..b6e9c8047a 100644 --- a/src/internet_identity/src/stats/event_stats/tests.rs +++ b/src/internet_identity/src/stats/event_stats/tests.rs @@ -637,6 +637,70 @@ 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( + EventData { + event: Event::PrepareDelegation(PrepareDelegationEvent { + ii_domain: ii_domain.clone(), + frontend: "https://example3.com".to_string(), + session_duration_ns: to_ns(0), + }), + }, + 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>>>) { assert_eq!(