Reduce size of data in stable memory (#46)

Cargo.toml

@@ -16,7 +16,7 @@ ic-agent = "0.33.0"
 ic-cdk = "0.12.0"
 ic-cdk-timers = "0.6.0"
 ic_principal = "0.1.1"
-ic-stable-structures = "0.6.2"
+ic-stable-structures = "0.6.3"
 pocket-ic = "2.1.0"
 rand = "0.8.5"
 rmp-serde = "1.1.2"

rs/canister/api/can.did

@@ -1,7 +1,14 @@
+type AnonymizationInitConfig = record {
+  users : opt bool;
+  exclusions : opt vec text;
+  sources : opt bool;
+};
 type EventsArgs = record { start : nat64; length : nat64 };
 type EventsResponse = record {
+  latest_event_index_v2 : opt nat64;
   events : vec IndexedEvent;
   latest_event_index : opt nat64;
+  is_v2 : bool;
 };
 type IdempotentEvent = record {
   source : opt text;
@@ -21,6 +28,7 @@ type IndexedEvent = record {
 };
 type InitArgs = record {
   push_events_whitelist : vec principal;
+  anonymization_config : opt AnonymizationInitConfig;
   read_events_whitelist : vec principal;
 };
 type PushEventsArgs = record { events : vec IdempotentEvent };

rs/canister/api/src/queries/events.rs

@@ -12,4 +12,6 @@ pub struct EventsArgs {
 pub struct EventsResponse {
     pub events: Vec<IndexedEvent>,
     pub latest_event_index: Option<u64>,
+    pub latest_event_index_v2: Option<u64>,
+    pub is_v2: bool,
 }

rs/canister/impl/src/lifecycle/heartbeat.rs

@@ -0,0 +1,7 @@
+use crate::state;
+use ic_cdk::heartbeat;
+
+#[heartbeat]
+fn heartbeat() {
+    state::mutate(|s| s.migrate_events(10000));
+}

rs/canister/impl/src/lifecycle/mod.rs

@@ -1,3 +1,4 @@
+mod heartbeat;
 mod init;
 mod post_upgrade;
 mod pre_upgrade;

rs/canister/impl/src/memory.rs

@@ -6,6 +6,11 @@ use ic_stable_structures::{
 const UPGRADES: MemoryId = MemoryId::new(0);
 const EVENTS_INDEX: MemoryId = MemoryId::new(1);
 const EVENTS_DATA: MemoryId = MemoryId::new(2);
+const EVENTS_V2_INDEX: MemoryId = MemoryId::new(3);
+const EVENTS_V2_DATA: MemoryId = MemoryId::new(4);
+const STRING_TO_NUM_MAP: MemoryId = MemoryId::new(5);
+const NUM_TO_STRING_INDEX: MemoryId = MemoryId::new(6);
+const NUM_TO_STRING_DATA: MemoryId = MemoryId::new(7);
 
 pub type Memory = VirtualMemory<DefaultMemoryImpl>;
 
@@ -26,6 +31,26 @@ pub fn get_events_data_memory() -> Memory {
     get_memory(EVENTS_DATA)
 }
 
+pub fn get_events_v2_index_memory() -> Memory {
+    get_memory(EVENTS_V2_INDEX)
+}
+
+pub fn get_events_v2_data_memory() -> Memory {
+    get_memory(EVENTS_V2_DATA)
+}
+
+pub fn get_string_to_num_map_memory() -> Memory {
+    get_memory(STRING_TO_NUM_MAP)
+}
+
+pub fn get_num_to_string_index_memory() -> Memory {
+    get_memory(NUM_TO_STRING_INDEX)
+}
+
+pub fn get_num_to_string_data_memory() -> Memory {
+    get_memory(NUM_TO_STRING_DATA)
+}
+
 fn get_memory(id: MemoryId) -> Memory {
     MEMORY_MANAGER.with(|m| m.get(id))
 }

rs/canister/impl/src/model/events_v2.rs

@@ -0,0 +1,128 @@
+use crate::memory::{get_events_v2_data_memory, get_events_v2_index_memory, Memory};
+use crate::model::string_to_num_map::StringToNumMap;
+use candid::Deserialize;
+use event_store_canister::{IdempotentEvent, IndexedEvent, TimestampMillis};
+use ic_stable_structures::storable::Bound;
+use ic_stable_structures::{StableLog, Storable};
+use serde::Serialize;
+use std::borrow::Cow;
+
+#[derive(Serialize, Deserialize)]
+pub struct EventsV2 {
+    #[serde(skip, default = "init_events")]
+    events: StableLog<StorableEvent, Memory, Memory>,
+    #[serde(default)]
+    string_to_num_map: StringToNumMap,
+}
+
+impl EventsV2 {
+    pub fn get(&self, start: u64, length: u64) -> Vec<IndexedEvent> {
+        self.events
+            .iter()
+            .skip(start as usize)
+            .take(length as usize)
+            .map(|e| self.hydrate(e))
+            .collect()
+    }
+
+    pub fn push(&mut self, event: IdempotentEvent) {
+        let storable = self.convert_to_storable(event, self.events.len());
+
+        self.events.append(&storable).unwrap();
+    }
+
+    pub fn stats(&self) -> EventsStats {
+        EventsStats {
+            latest_event_index: self.events.len().checked_sub(1),
+        }
+    }
+
+    pub fn len(&self) -> u64 {
+        self.events.len()
+    }
+
+    fn convert_to_storable(&mut self, event: IdempotentEvent, index: u64) -> StorableEvent {
+        StorableEvent {
+            index,
+            name: self.string_to_num_map.convert_to_num(event.name),
+            timestamp: event.timestamp,
+            user: event.user.map(|u| self.string_to_num_map.convert_to_num(u)),
+            source: event
+                .source
+                .map(|s| self.string_to_num_map.convert_to_num(s)),
+            payload: event.payload,
+        }
+    }
+
+    fn hydrate(&self, event: StorableEvent) -> IndexedEvent {
+        IndexedEvent {
+            index: event.index,
+            name: self
+                .string_to_num_map
+                .convert_to_string(event.name)
+                .unwrap_or("unknown".to_string()),
+            timestamp: event.timestamp,
+            user: event
+                .user
+                .and_then(|u| self.string_to_num_map.convert_to_string(u)),
+            source: event
+                .source
+                .and_then(|s| self.string_to_num_map.convert_to_string(s)),
+            payload: event.payload,
+        }
+    }
+}
+
+impl Default for EventsV2 {
+    fn default() -> Self {
+        EventsV2 {
+            events: init_events(),
+            string_to_num_map: StringToNumMap::default(),
+        }
+    }
+}
+
+fn init_events() -> StableLog<StorableEvent, Memory, Memory> {
+    StableLog::init(get_events_v2_index_memory(), get_events_v2_data_memory()).unwrap()
+}
+
+pub struct EventsStats {
+    pub latest_event_index: Option<u64>,
+}
+
+#[derive(Serialize, Deserialize)]
+struct StorableEvent {
+    #[serde(rename = "i")]
+    index: u64,
+    #[serde(rename = "n")]
+    name: u32,
+    #[serde(rename = "t")]
+    timestamp: TimestampMillis,
+    #[serde(rename = "u", default, skip_serializing_if = "Option::is_none")]
+    user: Option<u32>,
+    #[serde(rename = "s", default, skip_serializing_if = "Option::is_none")]
+    source: Option<u32>,
+    #[serde(
+        rename = "p",
+        default,
+        skip_serializing_if = "is_empty_slice",
+        with = "serde_bytes"
+    )]
+    payload: Vec<u8>,
+}
+
+impl Storable for StorableEvent {
+    fn to_bytes(&self) -> Cow<[u8]> {
+        Cow::Owned(rmp_serde::to_vec_named(&self).unwrap())
+    }
+
+    fn from_bytes(bytes: Cow<[u8]>) -> Self {
+        rmp_serde::from_slice(bytes.as_ref()).unwrap()
+    }
+
+    const BOUND: Bound = Bound::Unbounded;
+}
+
+fn is_empty_slice<T>(vec: &[T]) -> bool {
+    vec.is_empty()
+}

rs/canister/impl/src/model/mod.rs

@@ -1,2 +1,4 @@
 pub mod events;
+pub mod events_v2;
 pub mod salt;
+mod string_to_num_map;

rs/canister/impl/src/model/string_to_num_map.rs

@@ -0,0 +1,52 @@
+use crate::memory::{
+    get_num_to_string_data_memory, get_num_to_string_index_memory, get_string_to_num_map_memory,
+    Memory,
+};
+use ic_stable_structures::{StableBTreeMap, StableLog};
+use serde::{Deserialize, Serialize};
+
+#[derive(Serialize, Deserialize)]
+pub struct StringToNumMap {
+    #[serde(skip, default = "init_string_to_num")]
+    string_to_num: StableBTreeMap<String, u32, Memory>,
+    #[serde(skip, default = "init_num_to_string")]
+    num_to_string: StableLog<String, Memory, Memory>,
+}
+
+impl StringToNumMap {
+    pub fn convert_to_num(&mut self, string: String) -> u32 {
+        if let Some(i) = self.string_to_num.get(&string) {
+            i
+        } else {
+            let i = self.num_to_string.len() as u32;
+            self.num_to_string.append(&string).unwrap();
+            self.string_to_num.insert(string, i);
+            i
+        }
+    }
+
+    pub fn convert_to_string(&self, num: u32) -> Option<String> {
+        self.num_to_string.get(num as u64)
+    }
+}
+
+impl Default for StringToNumMap {
+    fn default() -> Self {
+        StringToNumMap {
+            string_to_num: init_string_to_num(),
+            num_to_string: init_num_to_string(),
+        }
+    }
+}
+
+fn init_string_to_num() -> StableBTreeMap<String, u32, Memory> {
+    StableBTreeMap::init(get_string_to_num_map_memory())
+}
+
+fn init_num_to_string() -> StableLog<String, Memory, Memory> {
+    StableLog::init(
+        get_num_to_string_index_memory(),
+        get_num_to_string_data_memory(),
+    )
+    .unwrap()
+}

rs/canister/impl/src/queries/events.rs

@@ -6,12 +6,19 @@ use ic_cdk::query;
 #[query(guard = "caller_can_read_events")]
 fn events(args: EventsArgs) -> EventsResponse {
     state::read(|s| {
-        let events = s.events().get(args.start, args.length);
         let stats = s.events().stats();
+        let stats_v2 = s.events_v2().stats();
+        let (events, is_v2) = if stats.latest_event_index > stats_v2.latest_event_index {
+            (s.events().get(args.start, args.length), false)
+        } else {
+            (s.events_v2().get(args.start, args.length), true)
+        };
 
         EventsResponse {
             events,
             latest_event_index: stats.latest_event_index,
+            latest_event_index_v2: stats_v2.latest_event_index,
+            is_v2,
         }
     })
 }

rs/canister/impl/src/state.rs

@@ -1,5 +1,6 @@
 use crate::env;
 use crate::model::events::Events;
+use crate::model::events_v2::EventsV2;
 use crate::model::salt::Salt;
 use candid::Principal;
 use event_store_canister::{
@@ -21,6 +22,8 @@ pub struct State {
     push_events_whitelist: HashSet<Principal>,
     read_events_whitelist: HashSet<Principal>,
     events: Events,
+    #[serde(default)]
+    events_v2: EventsV2,
     event_deduper: EventDeduper,
     salt: Salt,
     anonymization_config: AnonymizationConfig,
@@ -61,6 +64,7 @@ impl State {
             push_events_whitelist,
             read_events_whitelist,
             events: Events::default(),
+            events_v2: EventsV2::default(),
             event_deduper: EventDeduper::default(),
             anonymization_config: anonymization_config.into(),
             salt: Salt::default(),
@@ -88,6 +92,10 @@ impl State {
         &self.events
     }
 
+    pub fn events_v2(&self) -> &EventsV2 {
+        &self.events_v2
+    }
+
     pub fn set_salt(&mut self, salt: [u8; 32]) {
         self.salt.set(salt);
     }
@@ -117,7 +125,24 @@ impl State {
             }
         }
 
-        self.events.push(event);
+        if self.events_v2.stats().latest_event_index >= self.events.stats().latest_event_index {
+            self.events_v2.push(event);
+        } else {
+            self.events.push(event);
+        }
+    }
+
+    pub fn migrate_events(&mut self, count: u32) {
+        for event in self.events.get(self.events_v2.len(), count as u64) {
+            self.events_v2.push(IdempotentEvent {
+                idempotency_key: 0,
+                name: event.name,
+                timestamp: event.timestamp,
+                user: event.user,
+                source: event.source,
+                payload: event.payload,
+            });
+        }
     }
 
     fn anonymize(&self, value: &str) -> String {

rs/integration_tests/src/lib.rs

@@ -64,7 +64,8 @@ fn push_then_read_events_succeeds() {
     assert_eq!(read_response.events.len(), 5);
     assert_eq!(read_response.events.first().unwrap().index, 0);
     assert_eq!(read_response.events.last().unwrap().index, 4);
-    assert_eq!(read_response.latest_event_index, Some(9));
+    assert_eq!(read_response.latest_event_index, None);
+    assert_eq!(read_response.latest_event_index_v2, Some(9));
 }
 
 #[test_case(true, true)]