1 // Copyright (C) 2024 Huawei Device Co., Ltd. 2 // Licensed under the Apache License, Version 2.0 (the "License"); 3 // you may not use this file except in compliance with the License. 4 // You may obtain a copy of the License at 5 // 6 // http://www.apache.org/licenses/LICENSE-2.0 7 // 8 // Unless required by applicable law or agreed to in writing, software 9 // distributed under the License is distributed on an "AS IS" BASIS, 10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 11 // See the License for the specific language governing permissions and 12 // limitations under the License. 13 14 use std::collections::HashMap; 15 use std::io; 16 use std::sync::{Arc, Mutex, OnceLock, Weak}; 17 18 use request_utils::lru::LRUCache; 19 use request_utils::task_id::TaskId; 20 21 use super::data::{self, restore_files, FileCache, RamCache}; 22 use crate::data::MAX_CACHE_SIZE; 23 24 const DEFAULT_RAM_CACHE_SIZE: u64 = 1024 * 1024 * 20; 25 const DEFAULT_FILE_CACHE_SIZE: u64 = 1024 * 1024 * 100; 26 27 pub struct CacheManager { 28 pub(crate) rams: Mutex<LRUCache<TaskId, Arc<RamCache>>>, 29 pub(crate) backup_rams: Mutex<HashMap<TaskId, Arc<RamCache>>>, 30 pub(crate) files: Mutex<LRUCache<TaskId, FileCache>>, 31 32 pub(crate) update_from_file_once: 33 Mutex<HashMap<TaskId, Arc<OnceLock<io::Result<Weak<RamCache>>>>>>, 34 pub(crate) ram_handle: Mutex<data::ResourceManager>, 35 pub(crate) file_handle: Mutex<data::ResourceManager>, 36 } 37 38 impl CacheManager { new() -> Self39 pub fn new() -> Self { 40 Self { 41 rams: Mutex::new(LRUCache::new()), 42 files: Mutex::new(LRUCache::new()), 43 backup_rams: Mutex::new(HashMap::new()), 44 update_from_file_once: Mutex::new(HashMap::new()), 45 46 ram_handle: Mutex::new(data::ResourceManager::new(DEFAULT_RAM_CACHE_SIZE)), 47 file_handle: Mutex::new(data::ResourceManager::new(DEFAULT_FILE_CACHE_SIZE)), 48 } 49 } 50 set_ram_cache_size(&self, size: u64)51 pub fn set_ram_cache_size(&self, size: u64) { 52 self.ram_handle.lock().unwrap().change_total_size(size); 53 CacheManager::apply_cache(&self.ram_handle, &self.rams, |a| RamCache::task_id(a), 0); 54 } 55 set_file_cache_size(&self, size: u64)56 pub fn set_file_cache_size(&self, size: u64) { 57 self.file_handle.lock().unwrap().change_total_size(size); 58 CacheManager::apply_cache(&self.file_handle, &self.files, FileCache::task_id, 0); 59 } 60 restore_files(&'static self)61 pub fn restore_files(&'static self) { 62 for task_id in restore_files() { 63 let Some(file_cache) = FileCache::try_restore(task_id.clone(), self) else { 64 continue; 65 }; 66 self.files.lock().unwrap().insert(task_id, file_cache); 67 } 68 } 69 fetch(&'static self, task_id: &TaskId) -> Option<Arc<RamCache>>70 pub fn fetch(&'static self, task_id: &TaskId) -> Option<Arc<RamCache>> { 71 self.get_cache(task_id) 72 } 73 remove(&self, task_id: TaskId)74 pub fn remove(&self, task_id: TaskId) { 75 self.files.lock().unwrap().remove(&task_id); 76 self.backup_rams.lock().unwrap().remove(&task_id); 77 self.rams.lock().unwrap().remove(&task_id); 78 self.update_from_file_once.lock().unwrap().remove(&task_id); 79 } 80 contains(&self, task_id: &TaskId) -> bool81 pub fn contains(&self, task_id: &TaskId) -> bool { 82 self.files.lock().unwrap().contains_key(task_id) 83 || self.backup_rams.lock().unwrap().contains_key(task_id) 84 || self.rams.lock().unwrap().contains_key(task_id) 85 } 86 get_cache(&'static self, task_id: &TaskId) -> Option<Arc<RamCache>>87 pub(crate) fn get_cache(&'static self, task_id: &TaskId) -> Option<Arc<RamCache>> { 88 let res = self.rams.lock().unwrap().get(task_id).cloned(); 89 res.or_else(|| self.backup_rams.lock().unwrap().get(task_id).cloned()) 90 .or_else(|| self.update_ram_from_file(task_id)) 91 } 92 apply_cache<T>( handle: &Mutex<data::ResourceManager>, caches: &Mutex<LRUCache<TaskId, T>>, task_id: fn(&T) -> &TaskId, size: usize, ) -> bool93 pub(super) fn apply_cache<T>( 94 handle: &Mutex<data::ResourceManager>, 95 caches: &Mutex<LRUCache<TaskId, T>>, 96 task_id: fn(&T) -> &TaskId, 97 size: usize, 98 ) -> bool { 99 loop { 100 if size > MAX_CACHE_SIZE as usize { 101 return false; 102 } 103 if handle.lock().unwrap().apply_cache_size(size as u64) { 104 return true; 105 }; 106 107 match caches.lock().unwrap().pop() { 108 Some(cache) => { 109 info!("CacheManager release cache {}", task_id(&cache).brief()); 110 } 111 None => { 112 info!("CacheManager release cache failed"); 113 return false; 114 } 115 } 116 } 117 } 118 } 119 120 #[cfg(test)] 121 mod ut_manage { 122 include!("../tests/ut/ut_manage.rs"); 123 } 124