|
| 1 | +use crate::cachestore::{CacheItem, CacheStore, RocksCacheStore}; |
| 2 | +use crate::config::ConfigObj; |
| 3 | +use crate::metastore::{IdRow, MetaStoreFs}; |
| 4 | +use crate::CubeError; |
| 5 | +use async_trait::async_trait; |
| 6 | +use log::trace; |
| 7 | +use std::path::Path; |
| 8 | +use std::sync::Arc; |
| 9 | +use tokio::sync::watch::{Receiver, Sender}; |
| 10 | + |
| 11 | +pub enum LazyRocksCacheStoreState { |
| 12 | + FromRemote { |
| 13 | + path: String, |
| 14 | + metastore_fs: Arc<dyn MetaStoreFs>, |
| 15 | + config: Arc<dyn ConfigObj>, |
| 16 | + init_flag: Sender<bool>, |
| 17 | + }, |
| 18 | + Closed {}, |
| 19 | + Initialized { |
| 20 | + store: Arc<RocksCacheStore>, |
| 21 | + }, |
| 22 | +} |
| 23 | + |
| 24 | +pub struct LazyRocksCacheStore { |
| 25 | + init_signal: Option<Receiver<bool>>, |
| 26 | + state: tokio::sync::RwLock<LazyRocksCacheStoreState>, |
| 27 | +} |
| 28 | + |
| 29 | +impl LazyRocksCacheStore { |
| 30 | + pub async fn load_from_dump( |
| 31 | + path: &Path, |
| 32 | + dump_path: &Path, |
| 33 | + metastore_fs: Arc<dyn MetaStoreFs>, |
| 34 | + config: Arc<dyn ConfigObj>, |
| 35 | + ) -> Result<Arc<Self>, CubeError> { |
| 36 | + let store = RocksCacheStore::load_from_dump(path, dump_path, metastore_fs, config).await?; |
| 37 | + |
| 38 | + Ok(Arc::new(Self { |
| 39 | + init_signal: None, |
| 40 | + state: tokio::sync::RwLock::new(LazyRocksCacheStoreState::Initialized { store }), |
| 41 | + })) |
| 42 | + } |
| 43 | + |
| 44 | + pub async fn load_from_remote( |
| 45 | + path: &str, |
| 46 | + metastore_fs: Arc<dyn MetaStoreFs>, |
| 47 | + config: Arc<dyn ConfigObj>, |
| 48 | + ) -> Result<Arc<Self>, CubeError> { |
| 49 | + let (init_flag, init_signal) = tokio::sync::watch::channel::<bool>(false); |
| 50 | + |
| 51 | + Ok(Arc::new(Self { |
| 52 | + init_signal: Some(init_signal), |
| 53 | + state: tokio::sync::RwLock::new(LazyRocksCacheStoreState::FromRemote { |
| 54 | + path: path.to_string(), |
| 55 | + metastore_fs, |
| 56 | + config, |
| 57 | + init_flag, |
| 58 | + }), |
| 59 | + })) |
| 60 | + } |
| 61 | + |
| 62 | + async fn init(&self) -> Result<Arc<RocksCacheStore>, CubeError> { |
| 63 | + { |
| 64 | + let guard = self.state.read().await; |
| 65 | + match &*guard { |
| 66 | + LazyRocksCacheStoreState::FromRemote { .. } => {} |
| 67 | + LazyRocksCacheStoreState::Closed { .. } => { |
| 68 | + return Err(CubeError::internal( |
| 69 | + "Unable to initialize Cache Store on lazy call, it was closed".to_string(), |
| 70 | + )); |
| 71 | + } |
| 72 | + LazyRocksCacheStoreState::Initialized { store } => { |
| 73 | + return Ok(store.clone()); |
| 74 | + } |
| 75 | + } |
| 76 | + } |
| 77 | + |
| 78 | + let mut guard = self.state.write().await; |
| 79 | + match &*guard { |
| 80 | + LazyRocksCacheStoreState::FromRemote { |
| 81 | + path, |
| 82 | + metastore_fs, |
| 83 | + config, |
| 84 | + // receiver will be closed on drop |
| 85 | + init_flag: _, |
| 86 | + } => { |
| 87 | + let store = |
| 88 | + RocksCacheStore::load_from_remote(&path, metastore_fs.clone(), config.clone()) |
| 89 | + .await?; |
| 90 | + |
| 91 | + *guard = LazyRocksCacheStoreState::Initialized { |
| 92 | + store: store.clone(), |
| 93 | + }; |
| 94 | + |
| 95 | + Ok(store) |
| 96 | + } |
| 97 | + _ => Err(CubeError::internal( |
| 98 | + "Unable to initialize Cache Store on lazy call, unexpected state".to_string(), |
| 99 | + )), |
| 100 | + } |
| 101 | + } |
| 102 | + |
| 103 | + pub async fn wait_upload_loop(&self) { |
| 104 | + if let Some(init_signal) = &self.init_signal { |
| 105 | + let _ = init_signal.clone().changed().await; |
| 106 | + } |
| 107 | + |
| 108 | + let store = { |
| 109 | + let guard = self.state.read().await; |
| 110 | + if let LazyRocksCacheStoreState::Initialized { store } = &*guard { |
| 111 | + store.clone() |
| 112 | + } else { |
| 113 | + return (); |
| 114 | + } |
| 115 | + }; |
| 116 | + |
| 117 | + trace!("wait_upload_loop unblocked, Cache Store was initialized"); |
| 118 | + |
| 119 | + store.wait_upload_loop().await |
| 120 | + } |
| 121 | + |
| 122 | + pub async fn stop_processing_loops(&self) { |
| 123 | + let store = { |
| 124 | + let mut guard = self.state.write().await; |
| 125 | + match &*guard { |
| 126 | + LazyRocksCacheStoreState::Closed { .. } => { |
| 127 | + return (); |
| 128 | + } |
| 129 | + LazyRocksCacheStoreState::FromRemote { .. } => { |
| 130 | + *guard = LazyRocksCacheStoreState::Closed {}; |
| 131 | + |
| 132 | + return (); |
| 133 | + } |
| 134 | + LazyRocksCacheStoreState::Initialized { store } => { |
| 135 | + let store_to_move = store.clone(); |
| 136 | + |
| 137 | + *guard = LazyRocksCacheStoreState::Closed {}; |
| 138 | + |
| 139 | + store_to_move |
| 140 | + } |
| 141 | + } |
| 142 | + }; |
| 143 | + |
| 144 | + trace!("stop_processing_loops unblocked, Cache Store was initialized"); |
| 145 | + |
| 146 | + store.stop_processing_loops().await |
| 147 | + } |
| 148 | +} |
| 149 | + |
| 150 | +#[async_trait] |
| 151 | +impl CacheStore for LazyRocksCacheStore { |
| 152 | + async fn cache_all(&self) -> Result<Vec<IdRow<CacheItem>>, CubeError> { |
| 153 | + self.init().await?.cache_all().await |
| 154 | + } |
| 155 | + |
| 156 | + async fn cache_set( |
| 157 | + &self, |
| 158 | + item: CacheItem, |
| 159 | + update_if_not_exists: bool, |
| 160 | + ) -> Result<bool, CubeError> { |
| 161 | + self.init() |
| 162 | + .await? |
| 163 | + .cache_set(item, update_if_not_exists) |
| 164 | + .await |
| 165 | + } |
| 166 | + |
| 167 | + async fn cache_truncate(&self) -> Result<(), CubeError> { |
| 168 | + self.init().await?.cache_truncate().await |
| 169 | + } |
| 170 | + |
| 171 | + async fn cache_delete(&self, key: String) -> Result<(), CubeError> { |
| 172 | + self.init().await?.cache_delete(key).await |
| 173 | + } |
| 174 | + |
| 175 | + async fn cache_get(&self, key: String) -> Result<Option<IdRow<CacheItem>>, CubeError> { |
| 176 | + self.init().await?.cache_get(key).await |
| 177 | + } |
| 178 | + |
| 179 | + async fn cache_keys(&self, prefix: String) -> Result<Vec<IdRow<CacheItem>>, CubeError> { |
| 180 | + self.init().await?.cache_keys(prefix).await |
| 181 | + } |
| 182 | + |
| 183 | + async fn cache_incr(&self, path: String) -> Result<IdRow<CacheItem>, CubeError> { |
| 184 | + self.init().await?.cache_incr(path).await |
| 185 | + } |
| 186 | + |
| 187 | + async fn compaction(&self) -> Result<(), CubeError> { |
| 188 | + self.init().await?.compaction().await |
| 189 | + } |
| 190 | +} |
| 191 | + |
| 192 | +crate::di_service!(LazyRocksCacheStore, [CacheStore]); |
0 commit comments