feat(station): add cache layer to permission repository

core/station/impl/results.yml

@@ -1,25 +1,25 @@
 benches:
   find_external_canister_policies_are_below_query_limit:
     total:
-      instructions: 27299957
+      instructions: 27370495
       heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   list_external_canisters_with_all_statuses:
     total:
-      instructions: 3462083192
-      heap_increase: 19
+      instructions: 225638784
+      heap_increase: 0
       stable_memory_increase: 0
     scopes: {}
   repository_batch_insert_100_requests:
     total:
-      instructions: 490157820
+      instructions: 490537864
       heap_increase: 0
-      stable_memory_increase: 240
+      stable_memory_increase: 192
     scopes: {}
   repository_filter_all_request_ids_by_default_filters:
     total:
-      instructions: 466321334
+      instructions: 466672477
       heap_increase: 1
       stable_memory_increase: 0
     scopes: {}
@@ -31,13 +31,13 @@ benches:
     scopes: {}
   service_filter_all_requests_with_creation_time_filters:
     total:
-      instructions: 1151272168
+      instructions: 1127914406
       heap_increase: 0
       stable_memory_increase: 16
     scopes: {}
   service_filter_all_requests_with_default_filters:
     total:
-      instructions: 6015172659
+      instructions: 5837697222
       heap_increase: 3
       stable_memory_increase: 16
     scopes: {}

core/station/impl/src/controllers/system.rs

@@ -32,6 +32,7 @@ async fn initialize(input: Option<SystemInstall>) {
 pub async fn mock_init() {
     use crate::core::write_system_info;
     use crate::models::SystemInfo;
+    use crate::repositories::permission::PERMISSION_REPOSITORY;
     use candid::Principal;
 
     // Initialize the random number generator with a fixed seed to ensure deterministic
@@ -42,6 +43,9 @@ pub async fn mock_init() {
     let mut system = SystemInfo::default();
     system.set_upgrader_canister_id(Principal::from_slice(&[25; 29]));
 
+    // Initialize the permission cached entries for repositories.
+    PERMISSION_REPOSITORY.build_cache();
+
     write_system_info(system);
 }
 

core/station/impl/src/repositories/permission.rs

@@ -1,7 +1,8 @@
 use crate::{
+    core::ic_cdk::api::print,
     core::{with_memory_manager, Memory, PERMISSION_MEMORY_ID},
     models::{
-        permission::{Permission, PermissionKey},
+        permission::{Allow, Permission, PermissionKey},
         resource::{
             CallExternalCanisterResourceTarget, ExecutionMethodResourceTarget,
             ExternalCanisterResourceAction, Resource, ValidationMethodResourceTarget,
@@ -13,14 +14,16 @@ use candid::Principal;
 use ic_stable_structures::{memory_manager::VirtualMemory, StableBTreeMap};
 use lazy_static::lazy_static;
 use orbit_essentials::repository::Repository;
-use std::{cell::RefCell, sync::Arc};
+use std::{cell::RefCell, collections::HashMap, sync::Arc};
 
 thread_local! {
   static DB: RefCell<StableBTreeMap<PermissionKey, Permission, VirtualMemory<Memory>>> = with_memory_manager(|memory_manager| {
     RefCell::new(
       StableBTreeMap::init(memory_manager.get(PERMISSION_MEMORY_ID))
     )
-  })
+  });
+
+  static CACHE: RefCell<HashMap<Resource, Allow>> = RefCell::new(HashMap::new());
 }
 
 lazy_static! {
@@ -34,26 +37,118 @@ pub struct PermissionRepository {}
 
 impl Repository<PermissionKey, Permission> for PermissionRepository {
     fn list(&self) -> Vec<Permission> {
-        DB.with(|m| m.borrow().iter().map(|(_, v)| v).collect())
+        DB.with(|m| match self.use_only_cache() {
+            true => CACHE.with(|cache| {
+                cache
+                    .borrow()
+                    .iter()
+                    .map(|(resource, allow)| Permission {
+                        resource: resource.clone(),
+                        allow: allow.clone(),
+                    })
+                    .collect()
+            }),
+            false => m.borrow().iter().map(|(_, v)| v.clone()).collect(),
+        })
     }
 
     fn get(&self, key: &PermissionKey) -> Option<Permission> {
-        DB.with(|m| m.borrow().get(key))
+        DB.with(|m| {
+            let maybe_cache_hit = CACHE.with(|cache| {
+                cache.borrow().get(key).map(|allow| Permission {
+                    resource: key.clone(),
+                    allow: allow.clone(),
+                })
+            });
+
+            match self.use_only_cache() {
+                true => maybe_cache_hit,
+                false => maybe_cache_hit.or_else(|| m.borrow().get(key).clone()),
+            }
+        })
     }
 
     fn insert(&self, key: PermissionKey, value: Permission) -> Option<Permission> {
-        DB.with(|m| m.borrow_mut().insert(key, value.clone()))
+        DB.with(|m| {
+            let prev = m.borrow_mut().insert(key, value.clone());
+
+            // Update the cache and remove the first element if the cache is full.
+            CACHE.with(|cache| {
+                let mut cache = cache.borrow_mut();
+
+                cache.insert(value.resource.clone(), value.allow.clone());
+
+                if cache.len() >= PermissionRepository::MAX_CACHE_SIZE {
+                    let maybe_remove_key = cache.iter_mut().take(1).map(|(k, _)| k.clone()).next();
+                    if let Some(remove_key) = maybe_remove_key {
+                        cache.remove(&remove_key);
+                    }
+                }
+            });
+
+            prev
+        })
     }
 
     fn remove(&self, key: &PermissionKey) -> Option<Permission> {
-        DB.with(|m| m.borrow_mut().remove(key))
+        DB.with(|m| {
+            CACHE.with(|cache| cache.borrow_mut().remove(key));
+
+            m.borrow_mut().remove(key)
+        })
     }
 
     fn len(&self) -> usize {
         DB.with(|m| m.borrow().len()) as usize
     }
 }
 
+impl PermissionRepository {
+    /// Currently the cache uses around 0.35KiB per entry (Resource, Allow),
+    /// so the max cache size is around 108MiB.
+    ///
+    /// Moreover, it takes approximately 70million instructions to load each entry
+    /// to the cache, which means that rebuilding the cache from the repository
+    /// would take around 20B instructions.
+    ///
+    /// Since init/upgrade hooks can use up to 200B instructions, rebuilding
+    /// a cache in the worst case would take up to 10% of the available instructions.
+    pub const MAX_CACHE_SIZE: usize = 300_000;
+
+    /// Checks if every permission in the repository is in the cache.
+    fn use_only_cache(&self) -> bool {
+        self.len() <= Self::MAX_CACHE_SIZE
+    }
+
+    /// Builds the cache from the stable memory repository.
+    ///
+    /// This method should only be called during init or upgrade hooks to ensure that the cache is
+    /// up-to-date with the repository and that we have enough instructions to rebuild the cache.
+    pub fn build_cache(&self) {
+        if self.len() > Self::MAX_CACHE_SIZE {
+            print(format!(
+                "Only the first {} permissions will be added to the cache, the reposity has {} permissions.",
+                Self::MAX_CACHE_SIZE,
+                PERMISSION_REPOSITORY.len(),
+            ));
+
+            return;
+        }
+
+        CACHE.with(|cache| {
+            cache.borrow_mut().clear();
+
+            DB.with(|db| {
+                for (_, permission) in db.borrow().iter().take(Self::MAX_CACHE_SIZE) {
+                    cache
+                        .borrow_mut()
+                        .insert(permission.resource.clone(), permission.allow.clone());
+                }
+            });
+        });
+    }
+}
+
 impl PermissionRepository {
     /// Finds all permissions that are available for a given external canister.
     pub fn find_external_canister_call_permissions(
@@ -238,4 +333,75 @@ mod tests {
             ));
         }
     }
+
+    #[test]
+    fn insert_updates_cache() {
+        let repository = &PERMISSION_REPOSITORY;
+        let policy = mock_permission();
+
+        assert!(repository.get(&policy.key()).is_none());
+
+        repository.insert(policy.key(), policy.clone());
+
+        assert!(repository.get(&policy.key()).is_some());
+        assert!(CACHE.with(|cache| cache.borrow().contains_key(&policy.resource)));
+    }
+
+    #[test]
+    fn remove_updates_cache() {
+        let repository = &PERMISSION_REPOSITORY;
+        let policy = mock_permission();
+
+        assert!(repository.get(&policy.key()).is_none());
+
+        repository.insert(policy.key(), policy.clone());
+
+        assert!(repository.get(&policy.key()).is_some());
+        assert!(CACHE.with(|cache| cache.borrow().contains_key(&policy.resource)));
+
+        repository.remove(&policy.key());
+
+        assert!(repository.get(&policy.key()).is_none());
+        assert!(!CACHE.with(|cache| cache.borrow().contains_key(&policy.resource)));
+    }
+
+    #[test]
+    fn get_uses_cache() {
+        let repository = &PERMISSION_REPOSITORY;
+        let policy = mock_permission();
+
+        assert!(repository.get(&policy.key()).is_none());
+
+        repository.insert(policy.key(), policy.clone());
+
+        // Clear the stable repository to ensure that the cache is used.
+        DB.with(|db| db.borrow_mut().remove(&policy.key()));
+
+        assert!(repository.get(&policy.key()).is_some());
+    }
+
+    #[test]
+    fn list_uses_cache_when_all_entries_are_loaded() {
+        let repository = &PERMISSION_REPOSITORY;
+        let permissions = (0..10).map(|_| mock_permission()).collect::<Vec<_>>();
+
+        for permission in permissions.iter() {
+            repository.insert(permission.key(), permission.clone());
+        }
+
+        // Clear the stable repository to ensure that the cache is used.
+        DB.with(|db| {
+            let mut db = db.borrow_mut();
+
+            for permission in permissions.iter() {
+                db.remove(&permission.key());
+            }
+
+            assert_eq!(db.len(), 0);
+        });
+
+        // even though the stable repository is empty, the cache should still be in-place which
+        // would allow the list method to return the permissions.
+        assert_eq!(repository.list().len(), permissions.len());
+    }
 }

core/station/impl/src/services/system.rs

@@ -13,7 +13,7 @@ use crate::{
         system::{DisasterRecoveryCommittee, SystemInfo, SystemState},
         ManageSystemInfoOperationInput, RequestId, RequestKey, RequestStatus,
     },
-    repositories::{RequestRepository, REQUEST_REPOSITORY},
+    repositories::{permission::PERMISSION_REPOSITORY, RequestRepository, REQUEST_REPOSITORY},
 };
 use candid::Principal;
 use lazy_static::lazy_static;
@@ -227,6 +227,15 @@ impl SystemService {
         };
     }
 
+    /// Initializes the cache of the canister data.
+    ///
+    /// Must only be called within a canister init or post_upgrade call.
+    fn init_cache(&self) {
+        // Initializes the cache of the permission repository,
+        // using at most 20B instructions.
+        PERMISSION_REPOSITORY.build_cache();
+    }
+
     /// Initializes the canister with the given owners and settings.
     ///
     /// Must only be called within a canister init call.
@@ -252,6 +261,9 @@ impl SystemService {
         // sets the name of the canister
         system_info.set_name(input.name.clone());
 
+        // initializes the cache of the canister data, must happen during the same call as the init
+        self.init_cache();
+
         // Handles the post init process in a one-off timer to allow for inter canister calls,
         // this adds the default canister configurations, deploys the station upgrader and makes sure
         // there are no unintended controllers of the canister.
@@ -264,6 +276,9 @@ impl SystemService {
     ///
     /// Must only be called within a canister post_upgrade call.
     pub async fn upgrade_canister(&self, input: Option<SystemUpgrade>) -> ServiceResult<()> {
+        // initializes the cache of the canister data, must happen during the same call as the upgrade
+        self.init_cache();
+
         // recompute all metrics to make sure they are up to date, only gauges are recomputed
         // since they are the only ones that can change over time.
         recompute_metrics();