Skip to main content

mj_controller/
worker_client.rs

1//! Controller-side client for a session relay's JSON-lines proxy.
2
3use std::collections::{BTreeMap, BTreeSet, VecDeque};
4use std::path::Path;
5use std::process::Stdio;
6use std::sync::atomic::{AtomicBool, Ordering};
7use std::sync::{Arc, PoisonError};
8use std::time::{Duration, Instant};
9
10use anyhow::{Context, Result, anyhow, bail};
11use base64::Engine as _;
12use base64::engine::general_purpose::STANDARD as BASE64;
13use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWriteExt, BufReader};
14use tokio::process::{Child, ChildStdin, ChildStdout, Command};
15use tokio::sync::{mpsc, watch};
16
17use crate::targets::{
18    CommandSpec, SSH_MASTER_OPEN_TIMEOUT, SSH_RETRY_ATTEMPTS, SshAdmission, SshPermit, SshRefusal,
19    SshSessionLease, ssh_refusal,
20};
21use mj_core::config::harness_authentication_marker;
22use mj_core::credentials::{
23    CredentialSnapshot, CredentialSyncAction, CredentialSyncHandle, CredentialSyncOutcome,
24    CredentialSyncResult, CredentialSyncTarget, SYNC_INTERVAL, SyncAction, SyncTrigger, enqueue,
25    profiles_with_targets, read_credential_file, reconcile, validate_credential_payload,
26    write_credential_file,
27};
28use mj_core::elicitation::ElicitationResponse;
29use mj_core::relay::{
30    MAX_FRAME_BYTES, RELAY_EVENT_GENESIS_DIGEST, RELAY_MIN_PROTOCOL_VERSION,
31    RELAY_PROTOCOL_VERSION, RelayCommand, RelayCursor, RelayErrorCode, RelayEvent,
32    RelayOperationalState, RelayProtocolError, RelayRequest, RelayRequestEnvelope,
33    RelayResponseBody, RelayResponseEnvelope, RelayResponsePayload, RelayVersionRange,
34    ReviewerRequest, validate_relay_event,
35};
36
37pub use mj_client::session::{RelayAttachment, StartedReviewer};
38use mj_core::worker_launch::ReviewerLaunchConfig;
39
40const RELAY_RPC_TIMEOUT: Duration = Duration::from_secs(15);
41const RELAY_SLOW_OPERATION_WARNING: Duration = Duration::from_secs(5);
42/// Starting a target-side proxy may page the full worker executable in and
43/// traverse a container runtime before the relay sees `hello`. That is worker
44/// startup latency, not an ordinary in-connection RPC.
45const RELAY_HANDSHAKE_TIMEOUT: Duration = Duration::from_secs(300);
46/// An attachment can decompress a transport-sized page from cold journal
47/// segments. It remains bounded by the relay frame budget, but cold or loaded
48/// storage needs a filesystem deadline rather than an in-memory RPC deadline.
49const RELAY_HISTORY_TIMEOUT: Duration = Duration::from_secs(900);
50/// Advancing an acknowledgement can durably prune a large relay journal. The
51/// worker performs that maintenance before replying, so it needs a deadline
52/// sized for filesystem work rather than ordinary relay bookkeeping.
53const RELAY_ACKNOWLEDGE_TIMEOUT: Duration = Duration::from_secs(300);
54/// Capturing a review delta runs Git over every workspace repository, which is
55/// filesystem work on a possibly large tree rather than relay bookkeeping.
56const REVIEW_CAPTURE_TIMEOUT: Duration = Duration::from_secs(300);
57const RELAY_PROXY_DETACH_GRACE: Duration = Duration::from_millis(500);
58const RELAY_PROXY_REAP_POLL: Duration = Duration::from_millis(10);
59
60/// How many trailing stderr lines a failed connect reports back to its caller.
61const RELAY_PROXY_STDERR_TAIL: usize = 10;
62
63/// The waits between connection attempts while the worker has not bound its
64/// control socket yet, 1.55 s in all. The daemon connects about 34 ms after
65/// starting a worker, and the worker binds its socket within about a second:
66/// on launch-r9 the next attempt, half a second later, got in every time
67/// (R9-2).
68const WORKER_SOCKET_RETRY_DELAYS: [Duration; 5] = [
69    Duration::from_millis(50),
70    Duration::from_millis(100),
71    Duration::from_millis(200),
72    Duration::from_millis(400),
73    Duration::from_millis(800),
74];
75
76/// The proxy's last [`RELAY_PROXY_STDERR_TAIL`] non-empty stderr lines, shared
77/// with whoever has to report them.
78///
79/// The drain publishes each line here as it reads it, rather than returning
80/// the whole tail when it finishes. A failed connect has to bound how long it
81/// waits for the drain, because a proxy that leaves a grandchild holding
82/// stderr never reaches EOF. Reading the tail from here means that bound costs
83/// only the lines not yet read, instead of discarding every line already
84/// collected.
85type ProxyStderrTail = Arc<std::sync::Mutex<VecDeque<String>>>;
86
87/// Forward a relay proxy's stderr to the log, one line at a time, until the
88/// child closes it, keeping the tail in `tail`. Reporting rather than dropping
89/// keeps connect failures diagnosable now that the controller no longer shares
90/// its terminal, and lets a failed connect put the proxy's own complaint in
91/// the error the caller sees rather than only in the log.
92///
93/// Until hello completes, lines go to debug level: a failed connect reports
94/// its tail once, at the level the SSH refusal classifier gives it, so a
95/// routine MaxSessions refusal is not a warning (R7-1). Once the connection
96/// is up, anything the proxy says is a warning.
97async fn drain_proxy_stderr(
98    errors: tokio::process::ChildStderr,
99    purpose: String,
100    session_id: String,
101    tail: ProxyStderrTail,
102    handshake_done: Arc<AtomicBool>,
103) {
104    let mut lines = BufReader::new(errors).lines();
105    loop {
106        match lines.next_line().await {
107            Ok(Some(line)) if line.trim().is_empty() => continue,
108            Ok(Some(line)) => {
109                if handshake_done.load(Ordering::Acquire) {
110                    tracing::warn!(%session_id, %purpose, %line, "relay proxy stderr");
111                } else {
112                    tracing::debug!(%session_id, %purpose, %line, "relay proxy stderr");
113                }
114                let mut tail = tail.lock().unwrap_or_else(PoisonError::into_inner);
115                if tail.len() == RELAY_PROXY_STDERR_TAIL {
116                    tail.pop_front();
117                }
118                tail.push_back(line);
119            }
120            Ok(None) => return,
121            Err(error) => {
122                tracing::warn!(%session_id, %purpose, %error, "read relay proxy stderr");
123                return;
124            }
125        }
126    }
127}
128
129mod errors;
130pub use errors::*;
131mod connect;
132mod exchange;
133mod relay;
134mod reviewer;
135mod transport;
136use transport::*;
137mod credential_sync;
138pub use credential_sync::*;
139
140/// Controller-side connection to the durable ACP relay protocol.
141///
142/// This type does not construct transcript state or request unbounded history.
143/// Callers persist bounded attachment pages, then acknowledge only a frontier
144/// that is already durable locally.
145pub struct RelayClient {
146    child: Option<Child>,
147    input: Option<ChildStdin>,
148    output: BufReader<ChildStdout>,
149    request_timeout: Duration,
150    /// Why this connection can no longer be used, once a call gave up on a
151    /// reply that is still in flight. See [`RelayClient::exchange`].
152    abandoned: Option<String>,
153    next_request: u64,
154    connection_nonce: u64,
155    protocol_version: u32,
156    session_id: String,
157    relay_version: String,
158    /// Content address of the executable the worker is running, as reported in
159    /// hello. `None` from a worker built before the field existed.
160    worker_build: Option<String>,
161    latest_ordinal: u64,
162    latest_digest: String,
163    /// The shared-connection session the proxy runs on. Unlike the admission
164    /// permit, which is released once hello completes, the session is in use
165    /// for as long as the proxy runs, so the lease lives as long as `child`.
166    ssh_session: Option<SshSessionLease>,
167}
168
169#[cfg(test)]
170mod tests;