diff --git a/core/station/impl/results.yml b/core/station/impl/results.yml index 5f056f69a..ada410c5f 100644 --- a/core/station/impl/results.yml +++ b/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: {} diff --git a/core/station/impl/src/controllers/system.rs b/core/station/impl/src/controllers/system.rs index 60171a27e..9d716269b 100644 --- a/core/station/impl/src/controllers/system.rs +++ b/core/station/impl/src/controllers/system.rs @@ -32,6 +32,7 @@ async fn initialize(input: Option) { 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); } diff --git a/core/station/impl/src/repositories/permission.rs b/core/station/impl/src/repositories/permission.rs index f1be8e884..987b21fd8 100644 --- a/core/station/impl/src/repositories/permission.rs +++ b/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>> = with_memory_manager(|memory_manager| { RefCell::new( StableBTreeMap::init(memory_manager.get(PERMISSION_MEMORY_ID)) ) - }) + }); + + static CACHE: RefCell> = RefCell::new(HashMap::new()); } lazy_static! { @@ -34,19 +37,65 @@ pub struct PermissionRepository {} impl Repository for PermissionRepository { fn list(&self) -> Vec { - 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 { - 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 { - 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 { - 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 { @@ -54,6 +103,50 @@ impl Repository for PermissionRepository { } } +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(), + )); + } + + 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 +331,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::>(); + + 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()); + } } diff --git a/core/station/impl/src/services/system.rs b/core/station/impl/src/services/system.rs index 629b4fc5f..cdf8937b9 100644 --- a/core/station/impl/src/services/system.rs +++ b/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) -> 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();