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

coop/
probe.rs
