1use std::path::{Path, PathBuf};
2use std::sync::{Mutex, MutexGuard};
3
4use super::env::Settings;
5use super::fit::{self, encode};
6use super::host::{host_name, pid_start};
7use super::rate::Samples;
8use super::record::{Pgpu, Progress, Record};
9use super::util::{self, Clock, atomic_write};
10
11enum WriteKind {
12 Now,
13 Progress,
14}
15
16struct Inner {
17 record: Record,
18 dir: Option<PathBuf>,
19 off: bool,
20 finished: bool,
21 last_progress_write: Option<u64>,
22 rate_override: Option<f64>,
23 samples: Samples,
24 clock: Clock,
25 file_name: String,
26}
27
28#[must_use = "call finish or fail; process::exit does not run drop"]
29pub struct Job {
30 inner: Mutex<Inner>,
31}
32
33pub(crate) fn start_with(
34 tool: &str,
35 app: Option<&str>,
36 label: &str,
37 settings: Settings,
38 clock: Clock,
39) -> Job {
40 let now = clock.wall_ms();
41 let pid = util::process_id();
42 let host = host_name();
43 let tool = fit::cap(tool.to_string());
44 let id = format!("{host}:{pid}:{now}");
45 let file_name = format!("{pid}-{now}.json");
46 let pgpu = settings.gpu_queue_pid.map(|pid| Pgpu {
47 pid,
48 as_name: tool.clone(),
49 });
50 let record = Record {
51 v: 1,
52 id,
53 host,
54 pid,
55 pid_start: pid_start(),
56 tool,
57 app: app.map(|app| fit::cap(app.to_string())),
58 label: fit::cap(label.to_string()),
59 detail: None,
60 invoker: settings.invoker.map(fit::cap),
61 parent: settings.parent.map(fit::cap),
62 state: "running".to_string(),
63 stage: None,
64 stages: Vec::new(),
65 stage_note: None,
66 progress: None,
67 rate: None,
68 pgpu,
69 device: None,
70 out: None,
71 started: now,
72 updated: now,
73 ended: None,
74 exit: None,
75 error: None,
76 };
77 let job = Job {
78 inner: Mutex::new(Inner {
79 record,
80 dir: settings.dir,
81 off: settings.off || !util::WRITES,
82 finished: false,
83 last_progress_write: None,
84 rate_override: None,
85 samples: Samples::default(),
86 clock,
87 file_name,
88 }),
89 };
90 job.commit(WriteKind::Now, now);
91 job
92}
93
94impl Job {
95 pub fn detail(self, detail: impl Into<String>) -> Self {
96 self.set_detail(detail.into());
97 self
98 }
99
100 pub fn device(self, device: impl Into<String>) -> Self {
101 self.set_device(device.into());
102 self
103 }
104
105 pub fn out(self, path: impl AsRef<Path>) -> Self {
106 self.set_out(path.as_ref());
107 self
108 }
109
110 pub fn stage(&self, name: &str, total: Option<u64>, unit: &str) {
111 let mut inner = self.lock();
112 if inner.finished {
113 return;
114 }
115 let now = inner.clock.wall_ms();
116 inner.record.stage = Some(fit::cap(name.to_string()));
117 inner.record.stage_note = None;
118 inner.record.progress = Some(Progress {
119 done: 0,
120 total,
121 unit: fit::cap(unit.to_string()),
122 });
123 inner.samples.clear();
124 inner.last_progress_write = None;
125 let rate = inner.shown_rate();
126 inner.record.rate = rate;
127 inner.commit(WriteKind::Now, now);
128 }
129
130 pub fn stages(&self, names: &[&str]) {
131 let mut inner = self.lock();
132 if inner.finished {
133 return;
134 }
135 let now = inner.clock.wall_ms();
136 inner.record.stages = names
137 .iter()
138 .take(fit::MAX_STAGES)
139 .map(|name| fit::cap((*name).to_string()))
140 .collect();
141 inner.commit(WriteKind::Now, now);
142 }
143
144 pub fn stage_note(&self, note: impl Into<String>) {
145 let mut inner = self.lock();
146 if inner.finished {
147 return;
148 }
149 let now = inner.clock.wall_ms();
150 inner.record.stage_note = Some(fit::cap(note.into()));
151 inner.commit(WriteKind::Now, now);
152 }
153
154 pub fn progress(&self, done: u64) {
155 let mut inner = self.lock();
156 if inner.finished || inner.record.progress.is_none() {
157 return;
158 }
159 let now = inner.clock.wall_ms();
160 if let Some(progress) = inner.record.progress.as_mut() {
161 progress.done = done;
162 }
163 let mono = inner.clock.mono_ms();
164 inner.samples.push(mono, done);
165 let rate = inner.shown_rate();
166 inner.record.rate = rate;
167 inner.commit(WriteKind::Progress, now);
168 }
169
170 pub fn pgpu(&self, pid: Option<u32>) {
171 let mut inner = self.lock();
172 if inner.finished {
173 return;
174 }
175 let now = inner.clock.wall_ms();
176 let tool = inner.record.tool.clone();
177 inner.record.pgpu = pid.map(|pid| Pgpu { pid, as_name: tool });
178 inner.commit(WriteKind::Now, now);
179 }
180
181 pub fn rate(&self, rate: Option<f64>) {
182 let mut inner = self.lock();
183 if inner.finished {
184 return;
185 }
186 let now = inner.clock.wall_ms();
187 inner.rate_override = rate.filter(|value| value.is_finite());
188 let rate = inner.shown_rate();
189 inner.record.rate = rate;
190 inner.commit(WriteKind::Now, now);
191 }
192
193 pub fn finish(&self, code: i32) {
194 self.terminal(code == 0, code, None);
195 }
196
197 pub fn fail(&self, code: i32, error: impl Into<String>) {
198 let code = if code == 0 { 1 } else { code };
199 self.terminal(false, code, Some(error.into()));
200 }
201
202 pub fn child_env(&self) -> Vec<(&'static str, String)> {
203 let inner = self.lock();
204 let mut pairs = vec![("PITO_JOB_ID", inner.record.id.clone())];
205 if let Some(invoker) = &inner.record.invoker {
206 pairs.push(("PITO_INVOKER", invoker.clone()));
207 }
208 pairs
209 }
210
211 fn set_detail(&self, detail: String) {
212 let mut inner = self.lock();
213 if inner.finished {
214 return;
215 }
216 let now = inner.clock.wall_ms();
217 inner.record.detail = Some(fit::cap(detail));
218 inner.commit(WriteKind::Now, now);
219 }
220
221 fn set_device(&self, device: String) {
222 let mut inner = self.lock();
223 if inner.finished {
224 return;
225 }
226 let now = inner.clock.wall_ms();
227 inner.record.device = Some(fit::cap(device));
228 inner.commit(WriteKind::Now, now);
229 }
230
231 fn set_out(&self, path: &Path) {
232 let mut inner = self.lock();
233 if inner.finished {
234 return;
235 }
236 let now = inner.clock.wall_ms();
237 inner.record.out = Some(fit::cap(absolute(path)));
238 inner.commit(WriteKind::Now, now);
239 }
240
241 fn terminal(&self, ok: bool, code: i32, error: Option<String>) {
242 let mut inner = self.lock();
243 if inner.finished {
244 return;
245 }
246 let now = inner.clock.wall_ms();
247 inner.record.state = if ok {
248 "done".to_string()
249 } else {
250 "failed".to_string()
251 };
252 inner.record.exit = Some(code);
253 inner.record.error = error.map(fit::cap);
254 inner.record.ended = Some(now);
255 inner.finished = true;
256 inner.commit(WriteKind::Now, now);
257 }
258
259 fn commit(&self, kind: WriteKind, now: u64) {
260 self.lock().commit(kind, now);
261 }
262
263 fn lock(&self) -> MutexGuard<'_, Inner> {
264 self.inner
265 .lock()
266 .unwrap_or_else(|poison| poison.into_inner())
267 }
268}
269
270impl Drop for Job {
271 fn drop(&mut self) {
272 self.fail(1, "dropped");
273 }
274}
275
276impl Inner {
277 fn shown_rate(&self) -> Option<f64> {
278 if self
279 .record
280 .progress
281 .as_ref()
282 .is_some_and(|progress| progress.unit == "seconds")
283 {
284 return None;
285 }
286 if let Some(rate) = self.rate_override.filter(|value| value.is_finite()) {
287 return Some(rate);
288 }
289 self.samples.rate()
290 }
291
292 fn commit(&mut self, kind: WriteKind, now: u64) {
293 if self.off {
294 return;
295 }
296 let Some(dir) = self.dir.clone() else {
297 return;
298 };
299 let progress = matches!(kind, WriteKind::Progress);
300 let mono = self.clock.mono_ms();
301 if progress && !progress_due(self.last_progress_write, mono) {
302 return;
303 }
304 let previous_write = self.last_progress_write;
305 let previous_updated = self.record.updated;
306 if progress {
307 self.last_progress_write = Some(mono);
308 }
309 self.record.updated = now;
310 let bytes = encode(&self.record);
311 if bytes.is_empty() || atomic_write(&dir, &self.file_name, &bytes).is_err() {
312 if progress {
313 self.last_progress_write = previous_write;
314 }
315 self.record.updated = previous_updated;
316 }
317 }
318}
319
320fn progress_due(last: Option<u64>, now: u64) -> bool {
321 match last {
322 Some(last) if now >= last => now - last >= 1_000,
323 _ => true,
324 }
325}
326
327fn absolute(path: &Path) -> String {
328 if path.is_absolute() {
329 return path.to_string_lossy().into_owned();
330 }
331 match std::env::current_dir() {
332 Ok(cwd) => cwd.join(path).to_string_lossy().into_owned(),
333 Err(_) => path.to_string_lossy().into_owned(),
334 }
335}