xref: /DADK/dadk-user/src/executor/cache.rs (revision 73779f3d0abacaf05aae9b67820e68f4bb9cf53f)
1 use std::{
2     path::PathBuf,
3     sync::{Arc, Once},
4 };
5 
6 use log::info;
7 
8 use crate::{
9     parser::{
10         task::{CodeSource, DADKTask, TaskType},
11         task_log::TaskLog,
12     },
13     scheduler::SchedEntity,
14     utils::lazy_init::Lazy,
15 };
16 
17 use super::ExecutorError;
18 
19 pub static CACHE_ROOT: Lazy<PathBuf> = Lazy::new();
20 
21 /// # 初始化缓存根目录
22 ///
23 /// ## 参数
24 ///
25 /// - `path` 缓存根目录的路径
26 pub fn cache_root_init(path: Option<PathBuf>) -> Result<(), ExecutorError> {
27     let cache_root: String;
28     if path.is_none() {
29         // 查询环境变量,是否有设置缓存根目录
30         let env = std::env::var("DADK_CACHE_ROOT");
31         if env.is_ok() {
32             cache_root = env.unwrap();
33         } else {
34             // 如果没有设置环境变量,则使用默认值
35             // 默认值为当前目录下的.cache目录
36             let cwd = std::env::current_dir().map_err(|e| ExecutorError::IoError(e.to_string()))?;
37             let cwd = cwd.to_str();
38 
39             if cwd.is_none() {
40                 return Err(ExecutorError::IoError(
41                     std::io::Error::new(
42                         std::io::ErrorKind::Other,
43                         "Current dir is not a valid unicode string",
44                     )
45                     .to_string(),
46                 ));
47             }
48             let cwd = cwd.unwrap();
49 
50             cache_root = format!("{}/dadk_cache", cwd);
51         }
52     } else {
53         // 如果有设置缓存根目录,则使用设置的值
54         let path = path.unwrap();
55         let x = path.to_str().ok_or(ExecutorError::IoError(
56             std::io::Error::new(
57                 std::io::ErrorKind::Other,
58                 "Cache root dir is not a valid unicode string",
59             )
60             .to_string(),
61         ))?;
62         cache_root = x.to_string();
63     }
64 
65     let cache_root = PathBuf::from(cache_root);
66 
67     // 如果缓存根目录不存在,则创建
68     if !cache_root.exists() {
69         info!("Cache root dir not exists, create it: {:?}", cache_root);
70         std::fs::create_dir_all(&cache_root).map_err(|e| ExecutorError::IoError(e.to_string()))?;
71     } else if !cache_root.is_dir() {
72         // 如果缓存根目录不是目录,则报错
73         return Err(ExecutorError::IoError(
74             std::io::Error::new(
75                 std::io::ErrorKind::NotADirectory,
76                 format!("Cache root dir is not a directory: {:?}", cache_root),
77             )
78             .to_string(),
79         ));
80     }
81 
82     // 初始化缓存根目录
83     static CACHE_ROOT_INIT_ONCE: Once = Once::new();
84     CACHE_ROOT_INIT_ONCE.call_once(|| CACHE_ROOT.init(cache_root));
85 
86     // 设置环境变量
87     std::env::set_var("DADK_CACHE_ROOT", CACHE_ROOT.get().to_str().unwrap());
88     info!("Cache root dir: {:?}", CACHE_ROOT.get());
89     return Ok(());
90 }
91 
92 #[derive(Debug, Clone, Copy)]
93 pub enum CacheDirType {
94     /// 构建缓存目录
95     Build,
96     /// 源码缓存目录
97     Source,
98     /// 每个任务执行数据缓存目录
99     TaskData,
100 }
101 
102 #[derive(Debug, Clone)]
103 pub struct CacheDir {
104     #[allow(dead_code)]
105     entity: Arc<SchedEntity>,
106     pub path: PathBuf,
107     pub cache_type: CacheDirType,
108 }
109 
110 impl CacheDir {
111     pub const DADK_BUILD_CACHE_DIR_ENV_KEY_PREFIX: &'static str = "DADK_BUILD_CACHE_DIR";
112     pub const DADK_SOURCE_CACHE_DIR_ENV_KEY_PREFIX: &'static str = "DADK_SOURCE_CACHE_DIR";
113     pub fn new(entity: Arc<SchedEntity>, cache_type: CacheDirType) -> Result<Self, ExecutorError> {
114         let task = entity.task();
115         let path = Self::get_path(&task, cache_type);
116 
117         let result = Self {
118             entity,
119             path,
120             cache_type,
121         };
122 
123         result.create()?;
124 
125         return Ok(result);
126     }
127 
128     fn get_path(task: &DADKTask, cache_type: CacheDirType) -> PathBuf {
129         let cache_root = CACHE_ROOT.get();
130         let name_version = task.name_version();
131         let cache_dir = match cache_type {
132             CacheDirType::Build => {
133                 format!("{}/build/{}", cache_root.to_str().unwrap(), name_version)
134             }
135             CacheDirType::Source => {
136                 format!("{}/source/{}", cache_root.to_str().unwrap(), name_version)
137             }
138             CacheDirType::TaskData => {
139                 format!(
140                     "{}/task_data/{}",
141                     cache_root.to_str().unwrap(),
142                     name_version
143                 )
144             }
145         };
146 
147         return PathBuf::from(cache_dir);
148     }
149 
150     pub fn build_dir(entity: Arc<SchedEntity>) -> Result<PathBuf, ExecutorError> {
151         return Ok(Self::new(entity.clone(), CacheDirType::Build)?.path);
152     }
153 
154     pub fn source_dir(entity: Arc<SchedEntity>) -> Result<PathBuf, ExecutorError> {
155         return Ok(Self::new(entity.clone(), CacheDirType::Source)?.path);
156     }
157 
158     pub fn build_dir_env_key(entity: &Arc<SchedEntity>) -> Result<String, ExecutorError> {
159         let name_version_env = entity.task().name_version_env();
160         return Ok(format!(
161             "{}_{}",
162             Self::DADK_BUILD_CACHE_DIR_ENV_KEY_PREFIX,
163             name_version_env
164         ));
165     }
166 
167     pub fn source_dir_env_key(entity: &Arc<SchedEntity>) -> Result<String, ExecutorError> {
168         let name_version_env = entity.task().name_version_env();
169         return Ok(format!(
170             "{}_{}",
171             Self::DADK_SOURCE_CACHE_DIR_ENV_KEY_PREFIX,
172             name_version_env
173         ));
174     }
175 
176     pub fn need_source_cache(entity: &Arc<SchedEntity>) -> bool {
177         let task_type = &entity.task().task_type;
178 
179         if let TaskType::BuildFromSource(cs) = task_type {
180             match cs {
181                 CodeSource::Git(_) | CodeSource::Archive(_) => {
182                     return true;
183                 }
184                 CodeSource::Local(_) => {
185                     return false;
186                 }
187             }
188         } else if let TaskType::InstallFromPrebuilt(ps) = task_type {
189             match ps {
190                 crate::parser::task::PrebuiltSource::Archive(_) => return false,
191                 crate::parser::task::PrebuiltSource::Local(_) => return false,
192             }
193         }
194         unimplemented!("Not fully implemented task type: {:?}", task_type);
195     }
196 
197     pub fn create(&self) -> Result<(), ExecutorError> {
198         if !self.path.exists() {
199             info!("Cache dir not exists, create it: {:?}", self.path);
200             std::fs::create_dir_all(&self.path)
201                 .map_err(|e| ExecutorError::IoError(e.to_string()))?;
202             info!("Cache dir: [{:?}] created.", self.path);
203         } else if !self.path.is_dir() {
204             // 如果路径类别不是目录,则报错
205             return Err(ExecutorError::IoError(
206                 std::io::Error::new(
207                     std::io::ErrorKind::NotADirectory,
208                     format!("Cache dir is not a directory: {:?}", self.path),
209                 )
210                 .to_string(),
211             ));
212         }
213 
214         return Ok(());
215     }
216 
217     /// 判断缓存目录是否为空
218     pub fn is_empty(&self) -> Result<bool, ExecutorError> {
219         let x = self
220             .path
221             .read_dir()
222             .map_err(|e| ExecutorError::IoError(e.to_string()))?;
223         for _ in x {
224             return Ok(false);
225         }
226 
227         return Ok(true);
228     }
229 
230     /// # 递归删除自身目录
231     /// 递归删除自身目录,如果目录不存在,则忽略
232     ///
233     /// 请注意,这会删除整个目录,包括目录下的所有文件和子目录
234     pub fn remove_self_recursive(&self) -> Result<(), ExecutorError> {
235         let path = &self.path;
236         if path.exists() {
237             std::fs::remove_dir_all(path).map_err(|e| ExecutorError::IoError(e.to_string()))?;
238         }
239         return Ok(());
240     }
241 }
242 
243 #[derive(Debug, Clone)]
244 pub struct TaskDataDir {
245     dir: CacheDir,
246 }
247 
248 impl TaskDataDir {
249     const TASK_LOG_FILE_NAME: &'static str = "task_log.toml";
250     pub fn new(entity: Arc<SchedEntity>) -> Result<Self, ExecutorError> {
251         let dir = CacheDir::new(entity.clone(), CacheDirType::TaskData)?;
252         return Ok(Self { dir });
253     }
254 
255     /// # 获取任务日志
256     pub fn task_log(&self) -> TaskLog {
257         let path = self.dir.path.join(Self::TASK_LOG_FILE_NAME);
258         if path.exists() {
259             let content = std::fs::read_to_string(&path).unwrap();
260             let task_log: TaskLog = toml::from_str(&content).unwrap();
261             return task_log;
262         } else {
263             return TaskLog::new();
264         }
265     }
266 
267     /// # 设置任务日志
268     pub fn save_task_log(&self, task_log: &TaskLog) -> Result<(), ExecutorError> {
269         let path = self.dir.path.join(Self::TASK_LOG_FILE_NAME);
270         let content = toml::to_string(task_log).unwrap();
271         std::fs::write(&path, content).map_err(|e| ExecutorError::IoError(e.to_string()))?;
272         return Ok(());
273     }
274 }
275