coop/
probe.rs

1use std::time::Duration;
2
3use anyhow::{Context, Result, bail};
4
5use crate::config::Host;
6use crate::transport::Transport;
7use crate::wrapper::{JobId, state_dir};
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub enum State {
11    Running,
12    Done(i32),
13    Orphan,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct Probe {
18    pub state: State,
19    pub runtime_secs: Option<u64>,
20    pub log_size: u64,
21    pub bytes: Vec<u8>,
22}
23
24/// Ask for state only, fetching no log bytes.
25///
26/// A plain `wait` wants `rc`, not output, and shipping the log to discard it
27/// would hold the lock for the transfer. Expressed as its own type rather than
28/// a sentinel offset: passing `u64::MAX` overflowed the `+1` that `tail -c +N`
29/// needs and produced `tail: Invalid argument`, which surfaced as a spurious
30/// "lost contact while waiting" on every `--wait --no-tail`.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum From {
33    Offset(u64),
34    StateOnly,
35}
36
37impl From {
38    /// The `tail -c +N` argument, or `None` when no bytes are wanted.
39    fn tail_arg(self) -> Option<u64> {
40        match self {
41            // `tail -c +N` is 1-based, so byte offset 0 is `+1`. Saturating at
42            // the top keeps a nonsensical offset from wrapping into a valid
43            // one; it reads the last byte instead of the whole file.
44            From::Offset(n) => Some(n.saturating_add(1)),
45            From::StateOnly => None,
46        }
47    }
48}
49
50pub fn probe(t: &dyn Transport, host: &Host, id: &JobId, from: impl Into<From>) -> Result<Probe> {
51    let dir = state_dir(id);
52    let from = from.into();
53    let script = format!(
54        "d={dir}; rc=$(cat $d/rc 2>/dev/null); printf 'rc=%s\\n' \"$rc\"; \
55         alive=$(tmux -L {} has-session -t coop-{id} 2>/dev/null && echo 1 || echo 0); \
56         printf 'alive=%s\\n' \"$alive\"; \
57         cmd_mtime=$(stat -c %Y $d/cmd 2>/dev/null || stat -f %m $d/cmd 2>/dev/null); \
58         if [ -n \"$rc\" ]; then end_mtime=$(stat -c %Y $d/rc 2>/dev/null || stat -f %m $d/rc 2>/dev/null); \
59         elif [ \"$alive\" = 1 ]; then end_mtime=$(date +%s); else end_mtime=; fi; \
60         if [ -n \"$cmd_mtime\" ] && [ -n \"$end_mtime\" ]; then runtime=$((end_mtime - cmd_mtime)); \
61           [ \"$runtime\" -lt 0 ] && runtime=0; else runtime=; fi; \
62         printf 'runtime=%s\\n' \"$runtime\"; \
63         printf 'size=%s\\n' \"$(wc -c < $d/log 2>/dev/null || echo 0)\"; \
64         printf 'bytes:\\n'; {}",
65        host.tmux_socket,
66        match from.tail_arg() {
67            Some(n) => format!("tail -c +{n} $d/log 2>/dev/null"),
68            None => "true".to_string(),
69        }
70    );
71    let mut output = t.run(host, &script)?;
72    if output.code != 0 {
73        bail!("probe failed: {}", output.stderr.trim());
74    }
75
76    const MARKER: &[u8] = b"bytes:\n";
77    let marker = output
78        .stdout
79        .windows(MARKER.len())
80        .position(|window| window == MARKER)
81        .context("invalid probe reply: missing bytes marker")?;
82    let bytes = output.stdout[marker + MARKER.len()..].to_vec();
83    output.stdout.truncate(marker);
84
85    let mut rc = None;
86    let mut alive = None;
87    let mut runtime_secs = None;
88    let mut log_size = None;
89    for line in output.text().lines() {
90        if let Some(value) = line.strip_prefix("rc=") {
91            if !value.is_empty() {
92                rc = Some(value.parse::<i32>().context("invalid rc in probe reply")?);
93            }
94        } else if let Some(value) = line.strip_prefix("alive=") {
95            alive = Some(value == "1");
96        } else if let Some(value) = line.strip_prefix("runtime=") {
97            if !value.is_empty() {
98                runtime_secs = Some(value.parse().context("invalid runtime in probe reply")?);
99            }
100        } else if let Some(value) = line.strip_prefix("size=") {
101            log_size = Some(
102                value
103                    .trim()
104                    .parse()
105                    .context("invalid size in probe reply")?,
106            );
107        }
108    }
109
110    let state = match rc {
111        Some(code) => State::Done(code),
112        None if alive.context("invalid probe reply: missing alive")? => State::Running,
113        None => State::Orphan,
114    };
115    Ok(Probe {
116        state,
117        runtime_secs,
118        log_size: log_size.context("invalid probe reply: missing size")?,
119        bytes,
120    })
121}
122
123pub fn next_interval(current: Duration, new_bytes: bool) -> Duration {
124    if new_bytes {
125        Duration::from_secs(1)
126    } else {
127        current.saturating_mul(2).min(Duration::from_secs(5))
128    }
129}
130
131impl core::convert::From<u64> for From {
132    fn from(offset: u64) -> Self {
133        From::Offset(offset)
134    }
135}
coop/probe.rs

coop/
probe.rs
