xet_client/chunk_cache/
cache_manager.rs

1use std::cell::RefCell;
2use std::collections::HashMap;
3use std::path::PathBuf;
4use std::sync::{Arc, LazyLock, Mutex, Weak};
5
6use xet_runtime::config::XetConfig;
7
8use super::error::ChunkCacheError;
9use super::{CacheConfig, ChunkCache, DiskCache};
10
11// single instance of CACHE_MANAGER not exposed to outside users that
12// dedupes cache instances based on configurations
13static CACHE_MANAGER: LazyLock<CacheManager> = LazyLock::new(CacheManager::new);
14
15/// get_cache attempts to return a cache given the provided config parameter
16pub fn get_cache(xet_config: &XetConfig, config: &CacheConfig) -> Result<Arc<dyn ChunkCache>, ChunkCacheError> {
17    CACHE_MANAGER.get(xet_config, config)
18}
19
20struct CacheManager {
21    vals: Mutex<HashMap<PathBuf, RefCell<Weak<dyn ChunkCache>>>>,
22}
23
24impl CacheManager {
25    fn new() -> Self {
26        Self {
27            vals: Mutex::new(HashMap::new()),
28        }
29    }
30
31    /// get takes a CacheConfig and checks if there exists a valid `DiskCache` with a matching
32    /// cache_directory then it will return an Arc to that `DiskCache` instance. If it doesn't exist
33    /// or the `DiskCache` instance has been deallocated (CacheManager only holds a weak pointer)
34    /// then it creates a new instance based on the provided config.
35    fn get(&self, xet_config: &XetConfig, config: &CacheConfig) -> Result<Arc<dyn ChunkCache>, ChunkCacheError> {
36        let mut vals = self.vals.lock()?;
37        if let Some(v) = vals.get_mut(&config.cache_directory) {
38            let weak = v.borrow().clone();
39            // if upgrade from Weak to Arc is successful, returns the upgraded pointer
40            if let Some(value) = weak.upgrade() {
41                return Ok(value);
42            }
43            // since upgrading failed, creates a new DiskCache, replaces the weak pointer with a
44            // weak pointer to the new instance and then returns the Arc to the new cache instance
45            let result: Arc<dyn ChunkCache> = Arc::new(DiskCache::initialize(xet_config, config)?);
46            v.replace(Arc::downgrade(&result));
47            Ok(result)
48        } else {
49            // create a new Cache and insert weak pointer to managed map
50            let result: Arc<dyn ChunkCache> = Arc::new(DiskCache::initialize(xet_config, config)?);
51            vals.insert(config.cache_directory.clone(), RefCell::new(Arc::downgrade(&result)));
52            Ok(result)
53        }
54    }
55}
xet_client/chunk_cache/cache_manager.rs

xet_client/chunk_cache/
cache_manager.rs
