1
use std::collections::HashSet;
2
use std::path::PathBuf;
3
use std::sync::atomic::{
4
  AtomicBool,
5
  Ordering,
6
};
7
use std::sync::{
8
  Arc,
9
  Mutex,
10
};
11

            
12
use hashbrown::HashMap;
13
use indicatif::{
14
  MultiProgress,
15
  ProgressDrawTarget,
16
};
17
use serde::Serialize;
18

            
19
use crate::cache::CacheStore;
20
use crate::defaults::{
21
  default_ignore_errors,
22
  default_shell,
23
  default_verbose,
24
};
25
use crate::secrets::SecretConfig;
26

            
27
use super::{
28
  ActiveTasks,
29
  CompletedTasks,
30
  ContainerRuntime,
31
  Shell,
32
  TaskRoot,
33
};
34

            
35
/// Used to pass information to tasks
36
/// This use arc to allow for sharing of data between tasks
37
/// and allow parallel runs of tasks
38
#[derive(Clone)]
39
pub struct TaskContext {
40
  pub task_root: Arc<TaskRoot>,
41
  pub active_tasks: ActiveTasks,
42
  pub completed_tasks: CompletedTasks,
43
  pub multi: Arc<MultiProgress>,
44
  pub env_vars: HashMap<String, String>,
45
  pub task_outputs: Arc<Mutex<HashMap<String, String>>>,
46
  pub secret_config: Option<SecretConfig>,
47
  pub shell: Option<Arc<Shell>>,
48
  pub container_runtime: Option<ContainerRuntime>,
49
  pub ignore_errors: Option<bool>,
50
  pub verbose: Option<bool>,
51
  pub force: bool,
52
  pub json_events: bool,
53
  pub is_nested: bool,
54
  pub cache_store: Arc<Mutex<CacheStore>>,
55
  pub current_task_name: Option<String>,
56
  pub cancellation_requested: Arc<AtomicBool>,
57
}
58

            
59
impl TaskContext {
60
20
  pub fn empty() -> Self {
61
20
    let mp = MultiProgress::with_draw_target(ProgressDrawTarget::hidden());
62
20
    Self {
63
20
      task_root: Arc::new(TaskRoot::default()),
64
20
      active_tasks: Arc::new(Mutex::new(HashSet::new())),
65
20
      completed_tasks: Arc::new(Mutex::new(HashSet::new())),
66
20
      multi: Arc::new(mp),
67
20
      env_vars: HashMap::new(),
68
20
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
69
20
      secret_config: None,
70
20
      shell: None,
71
20
      container_runtime: None,
72
20
      ignore_errors: None,
73
20
      verbose: None,
74
20
      force: false,
75
20
      json_events: false,
76
20
      is_nested: false,
77
20
      cache_store: Arc::new(Mutex::new(CacheStore::default())),
78
20
      current_task_name: None,
79
20
      cancellation_requested: Arc::new(AtomicBool::new(false)),
80
20
    }
81
20
  }
82

            
83
6
  pub fn empty_with_root(task_root: Arc<TaskRoot>) -> Self {
84
6
    let mp = MultiProgress::with_draw_target(ProgressDrawTarget::hidden());
85
6
    Self {
86
6
      task_root: task_root.clone(),
87
6
      active_tasks: Arc::new(Mutex::new(HashSet::new())),
88
6
      completed_tasks: Arc::new(Mutex::new(HashSet::new())),
89
6
      multi: Arc::new(mp),
90
6
      env_vars: HashMap::new(),
91
6
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
92
6
      secret_config: None,
93
6
      shell: None,
94
6
      container_runtime: None,
95
6
      ignore_errors: None,
96
6
      verbose: None,
97
6
      force: false,
98
6
      json_events: false,
99
6
      is_nested: false,
100
6
      cache_store: Arc::new(Mutex::new(CacheStore::default())),
101
6
      current_task_name: None,
102
6
      cancellation_requested: Arc::new(AtomicBool::new(false)),
103
6
    }
104
6
  }
105

            
106
  pub fn new(task_root: Arc<TaskRoot>) -> Self {
107
    let cache_store = CacheStore::load_in_dir(&task_root.cache_base_dir()).unwrap_or_default();
108
    Self {
109
      task_root: task_root.clone(),
110
      active_tasks: Arc::new(Mutex::new(HashSet::new())),
111
      completed_tasks: Arc::new(Mutex::new(HashSet::new())),
112
      multi: Arc::new(MultiProgress::new()),
113
      env_vars: HashMap::new(),
114
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
115
      secret_config: None,
116
      shell: None,
117
      container_runtime: task_root.container_runtime.clone(),
118
      ignore_errors: None,
119
      verbose: None,
120
      force: false,
121
      json_events: false,
122
      is_nested: false,
123
      cache_store: Arc::new(Mutex::new(cache_store)),
124
      current_task_name: None,
125
      cancellation_requested: Arc::new(AtomicBool::new(false)),
126
    }
127
  }
128

            
129
80
  pub fn new_with_options(task_root: Arc<TaskRoot>, force: bool, json_events: bool) -> Self {
130
80
    let cache_store = CacheStore::load_in_dir(&task_root.cache_base_dir()).unwrap_or_default();
131
80
    let multi = if json_events {
132
8
      Arc::new(MultiProgress::with_draw_target(ProgressDrawTarget::hidden()))
133
    } else {
134
72
      Arc::new(MultiProgress::new())
135
    };
136
80
    Self {
137
80
      task_root: task_root.clone(),
138
80
      active_tasks: Arc::new(Mutex::new(HashSet::new())),
139
80
      completed_tasks: Arc::new(Mutex::new(HashSet::new())),
140
80
      multi,
141
80
      env_vars: HashMap::new(),
142
80
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
143
80
      secret_config: None,
144
80
      shell: None,
145
80
      container_runtime: task_root.container_runtime.clone(),
146
80
      ignore_errors: None,
147
80
      verbose: None,
148
80
      force,
149
80
      json_events,
150
80
      is_nested: false,
151
80
      cache_store: Arc::new(Mutex::new(cache_store)),
152
80
      current_task_name: None,
153
80
      cancellation_requested: Arc::new(AtomicBool::new(false)),
154
80
    }
155
80
  }
156

            
157
66
  pub fn from_context(context: &TaskContext) -> Self {
158
66
    Self {
159
66
      task_root: context.task_root.clone(),
160
66
      active_tasks: context.active_tasks.clone(),
161
66
      completed_tasks: context.completed_tasks.clone(),
162
66
      multi: context.multi.clone(),
163
66
      env_vars: context.env_vars.clone(),
164
66
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
165
66
      secret_config: context.secret_config.clone(),
166
66
      shell: context.shell.clone(),
167
66
      container_runtime: context.container_runtime.clone(),
168
66
      ignore_errors: context.ignore_errors,
169
66
      verbose: context.verbose,
170
66
      force: context.force,
171
66
      json_events: context.json_events,
172
66
      is_nested: true,
173
66
      cache_store: context.cache_store.clone(),
174
66
      current_task_name: context.current_task_name.clone(),
175
66
      cancellation_requested: context.cancellation_requested.clone(),
176
66
    }
177
66
  }
178

            
179
  pub fn from_context_with_args(context: &TaskContext, ignore_errors: bool, verbose: bool) -> Self {
180
    Self {
181
      task_root: context.task_root.clone(),
182
      active_tasks: context.active_tasks.clone(),
183
      completed_tasks: context.completed_tasks.clone(),
184
      multi: context.multi.clone(),
185
      env_vars: context.env_vars.clone(),
186
      task_outputs: Arc::new(Mutex::new(HashMap::new())),
187
      secret_config: context.secret_config.clone(),
188
      shell: context.shell.clone(),
189
      container_runtime: context.container_runtime.clone(),
190
      ignore_errors: Some(ignore_errors),
191
      verbose: Some(verbose),
192
      force: context.force,
193
      json_events: context.json_events,
194
      is_nested: true,
195
      cache_store: context.cache_store.clone(),
196
      current_task_name: context.current_task_name.clone(),
197
      cancellation_requested: context.cancellation_requested.clone(),
198
    }
199
  }
200

            
201
276
  pub fn extend_env_vars<I>(&mut self, iter: I)
202
276
  where
203
276
    I: IntoIterator<Item = (String, String)>,
204
  {
205
276
    self.env_vars.extend(iter);
206
276
  }
207

            
208
2
  pub fn set_shell(&mut self, shell: &Shell) {
209
2
    let shell = Arc::new(Shell::from_shell(shell));
210
2
    self.shell = Some(shell);
211
2
  }
212

            
213
68
  pub fn set_secret_config(&mut self, secret_config: SecretConfig) {
214
68
    self.secret_config = Some(secret_config);
215
68
  }
216

            
217
  pub fn set_container_runtime(&mut self, runtime: &ContainerRuntime) {
218
    self.container_runtime = Some(runtime.clone());
219
  }
220

            
221
2
  pub fn set_ignore_errors(&mut self, ignore_errors: bool) {
222
2
    self.ignore_errors = Some(ignore_errors);
223
2
  }
224

            
225
2
  pub fn set_verbose(&mut self, verbose: bool) {
226
2
    self.verbose = Some(verbose);
227
2
  }
228

            
229
4
  pub fn insert_task_output(&self, name: impl Into<String>, value: impl Into<String>) -> anyhow::Result<()> {
230
4
    let name = name.into();
231
4
    let mut outputs = self
232
4
      .task_outputs
233
4
      .lock()
234
4
      .map_err(|e| anyhow::anyhow!("Failed to lock task outputs - {}", e))?;
235
4
    if outputs.contains_key(&name) {
236
      anyhow::bail!("Task output already exists - {}", name);
237
4
    }
238
4
    outputs.insert(name, value.into());
239
4
    Ok(())
240
4
  }
241

            
242
4
  pub fn get_task_output(&self, name: &str) -> anyhow::Result<Option<String>> {
243
4
    let outputs = self
244
4
      .task_outputs
245
4
      .lock()
246
4
      .map_err(|e| anyhow::anyhow!("Failed to lock task outputs - {}", e))?;
247
4
    Ok(outputs.get(name).cloned())
248
4
  }
249

            
250
2
  pub fn has_task_output(&self, name: &str) -> anyhow::Result<bool> {
251
2
    let outputs = self
252
2
      .task_outputs
253
2
      .lock()
254
2
      .map_err(|e| anyhow::anyhow!("Failed to lock task outputs - {}", e))?;
255
2
    Ok(outputs.contains_key(name))
256
2
  }
257

            
258
74
  pub fn shell(&self) -> Arc<Shell> {
259
74
    self.shell.clone().unwrap_or_else(|| Arc::new(default_shell()))
260
74
  }
261

            
262
4
  pub fn ignore_errors(&self) -> bool {
263
4
    self.ignore_errors.unwrap_or(default_ignore_errors())
264
4
  }
265

            
266
4
  pub fn verbose(&self) -> bool {
267
4
    self.verbose.unwrap_or(default_verbose())
268
4
  }
269

            
270
90
  pub fn is_task_active(&self, task_name: &str) -> anyhow::Result<bool> {
271
90
    let active = self
272
90
      .active_tasks
273
90
      .lock()
274
90
      .map_err(|e| anyhow::anyhow!("Failed to lock active tasks - {}", e))?;
275
90
    Ok(active.contains(task_name))
276
90
  }
277

            
278
94
  pub fn is_task_completed(&self, task_name: &str) -> anyhow::Result<bool> {
279
94
    let completed = self
280
94
      .completed_tasks
281
94
      .lock()
282
94
      .map_err(|e| anyhow::anyhow!("Failed to lock completed tasks - {}", e))?;
283
94
    Ok(completed.contains(task_name))
284
94
  }
285

            
286
66
  pub fn mark_task_active(&self, task_name: &str) -> anyhow::Result<()> {
287
66
    let mut active = self
288
66
      .active_tasks
289
66
      .lock()
290
66
      .map_err(|e| anyhow::anyhow!("Failed to lock active tasks - {}", e))?;
291
66
    active.insert(task_name.to_string());
292
66
    Ok(())
293
66
  }
294

            
295
66
  pub fn unmark_task_active(&self, task_name: &str) -> anyhow::Result<()> {
296
66
    let mut active = self
297
66
      .active_tasks
298
66
      .lock()
299
66
      .map_err(|e| anyhow::anyhow!("Failed to lock active tasks - {}", e))?;
300
66
    active.remove(task_name);
301
66
    Ok(())
302
66
  }
303

            
304
66
  pub fn mark_task_complete(&self, task_name: &str) -> anyhow::Result<()> {
305
66
    let mut completed = self
306
66
      .completed_tasks
307
66
      .lock()
308
66
      .map_err(|e| anyhow::anyhow!("Failed to lock completed tasks - {}", e))?;
309
66
    completed.insert(task_name.to_string());
310
66
    Ok(())
311
66
  }
312

            
313
272
  pub fn emit_event<T: Serialize>(&self, value: &T) -> anyhow::Result<()> {
314
272
    if self.json_events {
315
32
      println!("{}", serde_json::to_string(value)?);
316
240
    }
317
272
    Ok(())
318
272
  }
319

            
320
66
  pub fn set_current_task_name(&mut self, task_name: &str) {
321
66
    self.current_task_name = Some(task_name.to_string());
322
66
  }
323

            
324
  pub fn request_cancellation(&self) {
325
    self.cancellation_requested.store(true, Ordering::SeqCst);
326
  }
327

            
328
4
  pub fn cancellation_requested(&self) -> bool {
329
4
    self.cancellation_requested.load(Ordering::SeqCst)
330
4
  }
331

            
332
4
  pub fn clear_cancellation(&self) {
333
4
    self.cancellation_requested.store(false, Ordering::SeqCst);
334
4
  }
335

            
336
  pub fn resolve_from_config(&self, value: &str) -> PathBuf {
337
    self.task_root.resolve_from_config(value)
338
  }
339
}
340

            
341
#[cfg(test)]
342
mod test {
343
  use super::*;
344

            
345
  #[test]
346
2
  fn test_task_context_1() -> anyhow::Result<()> {
347
    {
348
2
      let context = TaskContext::empty();
349
2
      let expected = if cfg!(windows) { "cmd" } else { "sh" };
350
2
      assert_eq!(context.shell().cmd(), expected.to_string());
351
2
      assert!(!context.ignore_errors());
352
2
      assert!(context.verbose());
353
    }
354

            
355
2
    Ok(())
356
2
  }
357

            
358
  #[test]
359
2
  fn test_task_context_2() -> anyhow::Result<()> {
360
    {
361
2
      let mut context = TaskContext::empty();
362
2
      context.set_shell(&Shell::String("bash".to_string()));
363
2
      assert_eq!(context.shell().cmd(), "bash".to_string());
364
    }
365

            
366
2
    Ok(())
367
2
  }
368

            
369
  #[test]
370
2
  fn test_task_context_3() -> anyhow::Result<()> {
371
    {
372
2
      let mut context = TaskContext::empty();
373
2
      context.extend_env_vars(vec![("key".to_string(), "value".to_string())]);
374
2
      assert_eq!(context.env_vars.get("key"), Some(&"value".to_string()));
375
    }
376

            
377
2
    Ok(())
378
2
  }
379

            
380
  #[test]
381
2
  fn test_task_context_4() -> anyhow::Result<()> {
382
    {
383
2
      let mut context = TaskContext::empty();
384
2
      context.set_ignore_errors(true);
385
2
      assert!(context.ignore_errors());
386
    }
387

            
388
2
    Ok(())
389
2
  }
390

            
391
  #[test]
392
2
  fn test_task_context_5() -> anyhow::Result<()> {
393
    {
394
2
      let mut context = TaskContext::empty();
395
2
      context.set_verbose(true);
396
2
      assert!(context.verbose());
397
    }
398

            
399
2
    Ok(())
400
2
  }
401

            
402
  #[test]
403
2
  fn test_task_context_outputs_are_stored() -> anyhow::Result<()> {
404
2
    let context = TaskContext::empty();
405
2
    context.insert_task_output("tag", "v1.0.0")?;
406
2
    assert_eq!(context.get_task_output("tag")?, Some("v1.0.0".to_string()));
407
2
    assert!(context.has_task_output("tag")?);
408
2
    Ok(())
409
2
  }
410
}