xref: /DADK/dadk-user/src/executor/cache.rs (revision d2ade6ef037f43d0342021730a0d8b3af7bc36c1)
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, path::abs_path},
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         abs_path(&PathBuf::from(cache_dir))
147     }
148 
149     pub fn build_dir(entity: Arc<SchedEntity>) -> Result<PathBuf, ExecutorError> {
150         return Ok(Self::new(entity.clone(), CacheDirType::Build)?.path);
151     }
152 
153     pub fn source_dir(entity: Arc<SchedEntity>) -> Result<PathBuf, ExecutorError> {
154         return Ok(Self::new(entity.clone(), CacheDirType::Source)?.path);
155     }
156 
157     pub fn build_dir_env_key(entity: &Arc<SchedEntity>) -> Result<String, ExecutorError> {
158         let name_version_env = entity.task().name_version_env();
159         return Ok(format!(
160             "{}_{}",
161             Self::DADK_BUILD_CACHE_DIR_ENV_KEY_PREFIX,
162             name_version_env
163         ));
164     }
165 
166     pub fn source_dir_env_key(entity: &Arc<SchedEntity>) -> Result<String, ExecutorError> {
167         let name_version_env = entity.task().name_version_env();
168         return Ok(format!(
169             "{}_{}",
170             Self::DADK_SOURCE_CACHE_DIR_ENV_KEY_PREFIX,
171             name_version_env
172         ));
173     }
174 
175     pub fn need_source_cache(entity: &Arc<SchedEntity>) -> bool {
176         let task_type = &entity.task().task_type;
177 
178         if let TaskType::BuildFromSource(cs) = task_type {
179             match cs {
180                 CodeSource::Git(_) | CodeSource::Archive(_) => {
181                     return true;
182                 }
183                 CodeSource::Local(_) => {
184                     return false;
185                 }
186             }
187         } else if let TaskType::InstallFromPrebuilt(ps) = task_type {
188             match ps {
189                 crate::parser::task::PrebuiltSource::Archive(_) => return false,
190                 crate::parser::task::PrebuiltSource::Local(_) => return false,
191             }
192         }
193         unimplemented!("Not fully implemented task type: {:?}", task_type);
194     }
195 
196     pub fn create(&self) -> Result<(), ExecutorError> {
197         if !self.path.exists() {
198             info!("Cache dir not exists, create it: {:?}", self.path);
199             std::fs::create_dir_all(&self.path)
200                 .map_err(|e| ExecutorError::IoError(e.to_string()))?;
201             info!("Cache dir: [{:?}] created.", self.path);
202         } else if !self.path.is_dir() {
203             // 如果路径类别不是目录,则报错
204             return Err(ExecutorError::IoError(
205                 std::io::Error::new(
206                     std::io::ErrorKind::NotADirectory,
207                     format!("Cache dir is not a directory: {:?}", self.path),
208                 )
209                 .to_string(),
210             ));
211         }
212 
213         return Ok(());
214     }
215 
216     /// 判断缓存目录是否为空
217     pub fn is_empty(&self) -> Result<bool, ExecutorError> {
218         let x = self
219             .path
220             .read_dir()
221             .map_err(|e| ExecutorError::IoError(e.to_string()))?;
222         for _ in x {
223             return Ok(false);
224         }
225 
226         return Ok(true);
227     }
228 
229     /// # 递归删除自身目录
230     /// 递归删除自身目录,如果目录不存在,则忽略
231     ///
232     /// 请注意,这会删除整个目录,包括目录下的所有文件和子目录
233     pub fn remove_self_recursive(&self) -> Result<(), ExecutorError> {
234         let path = &self.path;
235         if path.exists() {
236             std::fs::remove_dir_all(path).map_err(|e| ExecutorError::IoError(e.to_string()))?;
237         }
238         return Ok(());
239     }
240 }
241 
242 #[derive(Debug, Clone)]
243 pub struct TaskDataDir {
244     dir: CacheDir,
245 }
246 
247 impl TaskDataDir {
248     const TASK_LOG_FILE_NAME: &'static str = "task_log.toml";
249     pub fn new(entity: Arc<SchedEntity>) -> Result<Self, ExecutorError> {
250         let dir = CacheDir::new(entity.clone(), CacheDirType::TaskData)?;
251         return Ok(Self { dir });
252     }
253 
254     /// # 获取任务日志
255     pub fn task_log(&self) -> TaskLog {
256         let path = self.dir.path.join(Self::TASK_LOG_FILE_NAME);
257         if path.exists() {
258             let content = std::fs::read_to_string(&path).unwrap();
259             let task_log: TaskLog = toml::from_str(&content).unwrap();
260             return task_log;
261         } else {
262             return TaskLog::new();
263         }
264     }
265 
266     /// # 设置任务日志
267     pub fn save_task_log(&self, task_log: &TaskLog) -> Result<(), ExecutorError> {
268         let path = self.dir.path.join(Self::TASK_LOG_FILE_NAME);
269         let content = toml::to_string(task_log).unwrap();
270         std::fs::write(&path, content).map_err(|e| ExecutorError::IoError(e.to_string()))?;
271         return Ok(());
272     }
273 }
274