Move file metadata to stable memory

backend/canisters/storage_bucket/CHANGELOG.md

@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 
 - Reduce storage bucket memory usage ([#7103](https://github.com/open-chat-labs/open-chat/pull/7103))
+- Move file metadata to stable memory ([#7120](https://github.com/open-chat-labs/open-chat/pull/7120))
 
 ## [[2.0.1532](https://github.com/open-chat-labs/open-chat/releases/tag/v2.0.1532-storage_bucket)] - 2024-12-19
 

backend/canisters/storage_bucket/impl/Cargo.toml

@@ -24,10 +24,16 @@ rand = { workspace = true }
 serde = { workspace = true }
 serde_bytes = { workspace = true }
 stable_memory = { path = "../../../libraries/stable_memory" }
+stable_memory_map = { path = "../../../libraries/stable_memory_map" }
 storage_bucket_canister = { path = "../api" }
 storage_index_canister = { path = "../../storage_index/api" }
 storage_index_canister_c2c_client = { path = "../../storage_index/c2c_client" }
 timer_job_queues = { path = "../../../libraries/timer_job_queues" }
 tracing = { workspace = true }
 types = { path = "../../../libraries/types" }
 utils = { path = "../../../libraries/utils" }
+
+[dev-dependencies]
+proptest = { workspace = true }
+test-case = { workspace = true }
+test-strategy = { workspace = true }

backend/canisters/storage_bucket/impl/src/jobs/migrate_data_to_stable_map.rs

@@ -0,0 +1,39 @@
+use crate::mutate_state;
+use ic_cdk_timers::TimerId;
+use std::cell::Cell;
+use std::time::Duration;
+use tracing::info;
+
+thread_local! {
+    static TIMER_ID: Cell<Option<TimerId>> = Cell::default();
+}
+
+const INTERVAL: Duration = Duration::from_secs(60); // 1 minute
+
+pub fn start_job() {
+    let timer_id = ic_cdk_timers::set_timer_interval(INTERVAL, run);
+    TIMER_ID.set(Some(timer_id));
+}
+
+fn run() {
+    info!("'migrate_data_to_stable_map' job running");
+
+    mutate_state(|state| {
+        if state.data.files.migrate_files() {
+            state.data.files_migrated = true;
+
+            if state.data.files.migrate_reference_counts() {
+                state.data.file_reference_counts_migrated = true;
+
+                if state.data.files.migrate_accessors() {
+                    state.data.files_per_accessor_migrated = true;
+
+                    if let Some(timer_id) = TIMER_ID.take() {
+                        ic_cdk_timers::clear_timer(timer_id);
+                        info!("'migrate_data_to_stable_map' job completed");
+                    }
+                }
+            }
+        }
+    })
+}

backend/canisters/storage_bucket/impl/src/jobs/mod.rs

@@ -2,10 +2,12 @@ use crate::RuntimeState;
 
 pub mod calculate_freezing_limit;
 pub mod check_cycles_balance;
+mod migrate_data_to_stable_map;
 pub mod remove_expired_files;
 
 pub(crate) fn start(state: &RuntimeState) {
     calculate_freezing_limit::start_job();
     check_cycles_balance::start_job();
+    migrate_data_to_stable_map::start_job();
     remove_expired_files::start_job_if_required(state);
 }

backend/canisters/storage_bucket/impl/src/lib.rs

@@ -69,6 +69,9 @@ impl RuntimeState {
             index_sync_queue_length: self.data.index_event_sync_queue.len() as u32,
             freezing_limit: self.data.freezing_limit.value.unwrap_or_default(),
             stable_memory_sizes: memory::memory_sizes(),
+            files_migrated: self.data.files_migrated,
+            file_reference_counts_migrated: self.data.file_reference_counts_migrated,
+            files_per_accessor_migrated: self.data.files_per_accessor_migrated,
         }
     }
 }
@@ -83,11 +86,16 @@ struct Data {
     freezing_limit: Timestamped<Option<Cycles>>,
     rng_seed: [u8; 32],
     test_mode: bool,
+    #[serde(default)]
+    files_migrated: bool,
+    #[serde(default)]
+    file_reference_counts_migrated: bool,
+    #[serde(default)]
+    files_per_accessor_migrated: bool,
 }
 
 impl Data {
     pub fn new(storage_index_canister_id: CanisterId, now: TimestampMillis, test_mode: bool) -> Data {
-        #[allow(deprecated)]
         Data {
             storage_index_canister_id,
             users: Users::default(),
@@ -97,6 +105,9 @@ impl Data {
             freezing_limit: Timestamped::default(),
             rng_seed: [0; 32],
             test_mode,
+            files_migrated: true,
+            file_reference_counts_migrated: true,
+            files_per_accessor_migrated: true,
         }
     }
 
@@ -130,6 +141,9 @@ pub struct Metrics {
     pub index_sync_queue_length: u32,
     pub freezing_limit: Cycles,
     pub stable_memory_sizes: BTreeMap<u8, u64>,
+    pub files_migrated: bool,
+    pub file_reference_counts_migrated: bool,
+    pub files_per_accessor_migrated: bool,
 }
 
 pub fn calc_chunk_count(chunk_size: u32, total_size: u64) -> u32 {

backend/canisters/storage_bucket/impl/src/lifecycle/init.rs

@@ -1,4 +1,5 @@
 use crate::lifecycle::{init_env, init_state};
+use crate::memory::get_stable_memory_map_memory;
 use crate::Data;
 use canister_tracing_macros::trace;
 use ic_cdk::init;
@@ -10,6 +11,7 @@ use utils::env::Environment;
 #[trace]
 fn init(args: Args) {
     canister_logger::init(args.test_mode);
+    stable_memory_map::init(get_stable_memory_map_memory());
 
     let env = init_env([0; 32]);
 

backend/canisters/storage_bucket/impl/src/lifecycle/post_upgrade.rs

@@ -1,5 +1,5 @@
 use crate::lifecycle::{init_env, init_state};
-use crate::memory::get_upgrades_memory;
+use crate::memory::{get_stable_memory_map_memory, get_upgrades_memory};
 use crate::Data;
 use canister_logger::LogEntry;
 use canister_tracing_macros::trace;
@@ -11,6 +11,8 @@ use tracing::info;
 #[post_upgrade]
 #[trace]
 fn post_upgrade(args: Args) {
+    stable_memory_map::init(get_stable_memory_map_memory());
+
     let memory = get_upgrades_memory();
     let reader = get_reader(&memory);
 

backend/canisters/storage_bucket/impl/src/memory.rs

@@ -6,6 +6,7 @@ use std::collections::BTreeMap;
 
 const UPGRADES: MemoryId = MemoryId::new(0);
 const BLOBS: MemoryId = MemoryId::new(1);
+const STABLE_MEMORY_MAP: MemoryId = MemoryId::new(2);
 
 pub type Memory = VirtualMemory<DefaultMemoryImpl>;
 
@@ -22,8 +23,12 @@ pub fn get_blobs_memory() -> Memory {
     get_memory(BLOBS)
 }
 
+pub fn get_stable_memory_map_memory() -> Memory {
+    get_memory(STABLE_MEMORY_MAP)
+}
+
 pub fn memory_sizes() -> BTreeMap<u8, u64> {
-    (0u8..=1).map(|id| (id, get_memory(MemoryId::new(id)).size())).collect()
+    (0u8..=2).map(|id| (id, get_memory(MemoryId::new(id)).size())).collect()
 }
 
 fn get_memory(id: MemoryId) -> Memory {