Track registration rate and expose metrics (#2634)

* Track registration rate and expose metrics

This PR introduces the functionality of tracking the current and the
reference rate of new registrations.

In addition, it will also always calculate and expose the threshold rate
that would need to be crossed to trigger the captcha being activated.
This will be useful to later analyze how often the captcha was actually
shown.

Note: This is only the tracking / metrics. The captcha itslef is not yet
dynamic.

* Improved comments based on review feedback

* Doc indentation

* Cargo fmt

src/canister_tests/src/framework.rs

@@ -484,6 +484,11 @@ pub fn assert_metric(metrics: &str, metric_name: &str, expected: f64) {
     assert_eq!(value, expected, "metric {metric_name} does not match");
 }
 
+pub fn assert_metric_approx(metrics: &str, metric_name: &str, expected: f64, tolerance: f64) {
+    let (value, _) = parse_metric(metrics, metric_name);
+    assert!((value - expected).abs() <= tolerance, "metric {metric_name} is too far off: value={value}, expected={expected}, tolerance={tolerance}");
+}
+
 /// Asserts that the given metric is present in the metrics string and that it has the expected value
 /// across all the provided label values for the given label.
 pub fn assert_labelled_metric(

src/internet_identity/src/anchor_management/registration.rs

@@ -106,6 +106,7 @@ pub fn register(
                 anchor_number
             ))
         });
+        storage.registration_rates.new_registration()
     });
 
     // Save the 'temp_key' as a mean of authenticating for a short period of time, see

src/internet_identity/src/anchor_management/registration/registration_flow_v2.rs

@@ -153,8 +153,11 @@ fn create_identity(arg: &IdRegFinishArg) -> Result<IdentityNumber, IdRegFinishEr
 
     let identity_number = identity.anchor_number();
 
-    state::storage_borrow_mut(|s| s.write(identity))
-        .map_err(|err| IdRegFinishError::StorageError(err.to_string()))?;
+    state::storage_borrow_mut(|s| {
+        s.registration_rates.new_registration();
+        s.write(identity)
+    })
+    .map_err(|err| IdRegFinishError::StorageError(err.to_string()))?;
 
     let operation = Operation::RegisterAnchor {
         device: DeviceDataWithoutAlias::from(device),

src/internet_identity/src/http/metrics.rs

@@ -48,6 +48,24 @@ fn encode_metrics(w: &mut MetricsEncoder<Vec<u8>>) -> std::io::Result<()> {
             storage.event_aggregations.len() as f64,
             "Number of entries in the event_aggregations map.",
         )?;
+        if let Some(registration_rates) = storage.registration_rates.registration_rates() {
+            w.gauge_vec(
+                "internet_identity_registrations_per_second",
+                "Rate of new identity registrations on Internet Identity",
+            )?
+            .value(
+                &[("type", "reference_rate")],
+                registration_rates.reference_rate_per_second,
+            )?
+            .value(
+                &[("type", "current_rate")],
+                registration_rates.current_rate_per_second,
+            )?
+            .value(
+                &[("type", "captcha_threshold_rate")],
+                registration_rates.captcha_threshold_rate,
+            )?;
+        }
 
         let mut virtual_memory_stats_builder = w.gauge_vec(
             "internet_identity_virtual_memory_size_pages",

src/internet_identity/src/storage.rs

@@ -92,7 +92,9 @@ use ic_stable_structures::memory_manager::{MemoryId, MemoryManager, VirtualMemor
 use ic_stable_structures::reader::Reader;
 use ic_stable_structures::storable::Bound;
 use ic_stable_structures::writer::Writer;
-use ic_stable_structures::{Memory, RestrictedMemory, StableBTreeMap, StableCell, Storable};
+use ic_stable_structures::{
+    Memory, MinHeap, RestrictedMemory, StableBTreeMap, StableCell, Storable,
+};
 use internet_identity_interface::archive::types::BufferedEntry;
 
 use crate::stats::event_stats::AggregationKey;
@@ -102,10 +104,12 @@ use internet_identity_interface::internet_identity::types::*;
 use crate::state::PersistentState;
 use crate::storage::anchor::Anchor;
 use crate::storage::memory_wrapper::MemoryWrapper;
+use crate::storage::registration_rates::RegistrationRates;
 use crate::storage::storable_anchor::StorableAnchor;
 use crate::storage::storable_persistent_state::StorablePersistentState;
 
 pub mod anchor;
+pub mod registration_rates;
 
 /// module for the internal serialization format of anchors
 mod storable_anchor;
@@ -129,11 +133,17 @@ const ARCHIVE_BUFFER_MEMORY_INDEX: u8 = 1u8;
 const PERSISTENT_STATE_MEMORY_INDEX: u8 = 2u8;
 const EVENT_DATA_MEMORY_INDEX: u8 = 3u8;
 const STATS_AGGREGATIONS_MEMORY_INDEX: u8 = 4u8;
+const REGISTRATION_REFERENCE_RATE_MEMORY_INDEX: u8 = 5u8;
+const REGISTRATION_CURRENT_RATE_MEMORY_INDEX: u8 = 6u8;
 const ANCHOR_MEMORY_ID: MemoryId = MemoryId::new(ANCHOR_MEMORY_INDEX);
 const ARCHIVE_BUFFER_MEMORY_ID: MemoryId = MemoryId::new(ARCHIVE_BUFFER_MEMORY_INDEX);
 const PERSISTENT_STATE_MEMORY_ID: MemoryId = MemoryId::new(PERSISTENT_STATE_MEMORY_INDEX);
 const EVENT_DATA_MEMORY_ID: MemoryId = MemoryId::new(EVENT_DATA_MEMORY_INDEX);
 const STATS_AGGREGATIONS_MEMORY_ID: MemoryId = MemoryId::new(STATS_AGGREGATIONS_MEMORY_INDEX);
+const REGISTRATION_REFERENCE_RATE_MEMORY_ID: MemoryId =
+    MemoryId::new(REGISTRATION_REFERENCE_RATE_MEMORY_INDEX);
+const REGISTRATION_CURRENT_RATE_MEMORY_ID: MemoryId =
+    MemoryId::new(REGISTRATION_CURRENT_RATE_MEMORY_INDEX);
 // The bucket size 128 is relatively low, to avoid wasting memory when using
 // multiple virtual memories for smaller amounts of data.
 // This value results in 256 GB of total managed memory, which should be enough
@@ -186,6 +196,13 @@ pub struct Storage<M: Memory> {
     /// Memory wrapper used to report the size of the stats aggregation memory.
     event_aggregations_memory_wrapper: MemoryWrapper<ManagedMemory<M>>,
     pub event_aggregations: StableBTreeMap<AggregationKey, u64, ManagedMemory<M>>,
+    /// Registration rates tracked for the purpose of toggling the dynamic captcha (if configured)
+    /// This data is persisted as it potentially contains data collected over longer periods of time.
+    pub registration_rates: RegistrationRates<ManagedMemory<M>>,
+    /// Memory wrapper used to report the size of the current registration rate memory.
+    current_registration_rate_memory_wrapper: MemoryWrapper<ManagedMemory<M>>,
+    /// Memory wrapper used to report the size of the reference registration rate memory.
+    reference_registration_rate_memory_wrapper: MemoryWrapper<ManagedMemory<M>>,
 }
 
 #[repr(packed)]
@@ -243,10 +260,28 @@ impl<M: Memory + Clone> Storage<M> {
         let persistent_state_memory = memory_manager.get(PERSISTENT_STATE_MEMORY_ID);
         let event_data_memory = memory_manager.get(EVENT_DATA_MEMORY_ID);
         let stats_aggregations_memory = memory_manager.get(STATS_AGGREGATIONS_MEMORY_ID);
+        let registration_ref_rate_memory =
+            memory_manager.get(REGISTRATION_REFERENCE_RATE_MEMORY_ID);
+        let registration_current_rate_memory =
+            memory_manager.get(REGISTRATION_CURRENT_RATE_MEMORY_ID);
+
+        let registration_rates = RegistrationRates::new(
+            MinHeap::init(registration_ref_rate_memory.clone())
+                .expect("failed to initialize registration reference rate min heap"),
+            MinHeap::init(registration_current_rate_memory.clone())
+                .expect("failed to initialize registration current rate min heap"),
+        );
         Self {
             header,
             header_memory,
             anchor_memory,
+            registration_rates,
+            reference_registration_rate_memory_wrapper: MemoryWrapper::new(
+                registration_ref_rate_memory,
+            ),
+            current_registration_rate_memory_wrapper: MemoryWrapper::new(
+                registration_current_rate_memory,
+            ),
             archive_buffer_memory_wrapper: MemoryWrapper::new(archive_buffer_memory.clone()),
             archive_entries_buffer: StableBTreeMap::init(archive_buffer_memory),
             persistent_state_memory_wrapper: MemoryWrapper::new(persistent_state_memory.clone()),
@@ -514,6 +549,14 @@ impl<M: Memory + Clone> Storage<M> {
                 "event_aggregations".to_string(),
                 self.event_aggregations_memory_wrapper.size(),
             ),
+            (
+                "reference_registration_rate".to_string(),
+                self.reference_registration_rate_memory_wrapper.size(),
+            ),
+            (
+                "current_registration_rate".to_string(),
+                self.current_registration_rate_memory_wrapper.size(),
+            ),
         ])
     }
 }