mj_controller/hel_controller/
worker_restart.rs

1//! Replacing a session's worker process in place.
2//!
3//! Two things ask for this: a checkpoint whose ACP turn will not finish, and a
4//! session whose worker predates the controller now talking to it. Both stop
5//! the worker, install the binary this controller would provision, start it
6//! and reconnect, so the sequence lives here once and each caller supplies
7//! only what it tells the operator.
8
9use std::time::Duration;
10
11use anyhow::{Context, Result, bail};
12
13use crate::hel_session_manager::{SessionManagerControl, StandaloneSession};
14use hel::hel_targets::{self, CommandExecutor, CommandSpec};
15use hel::hel_worker::RelayExecutionState;
16
17use super::Controller;
18use super::readiness::{connect_started_worker_with_timeout, wait_for_native_session};
19use super::worker_binary::{
20    replace_installed_worker_binary, start_worker, stop_worker_after_target_recovery,
21    worker_binary_for, worker_probe_diagnosis,
22};
23
24/// How long a restarted worker has to recover its journal, bind `control.sock`
25/// and report an idle ACP session. Journal recovery over a long transcript
26/// runs before the socket exists, so this has to outlast it.
27const WORKER_RESTART_TIMEOUT: Duration = Duration::from_secs(300);
28
29/// How long a quiet session's upgrade waits for its actor and its lease.
30const UPGRADE_LEASE_TIMEOUT: Duration = Duration::from_secs(5);
31
32/// What one restart tells the operator at each step. The steps are identical;
33/// only the reason differs, and a diagnostic that named the wrong reason would
34/// send someone looking in the wrong place.
35pub(super) struct WorkerRestartMessages {
36    pub stop: &'static str,
37    pub replace: &'static str,
38    pub start: &'static str,
39    pub connect: &'static str,
40    pub project_memory: &'static str,
41    pub native_session: &'static str,
42}
43
44/// A wedged ACP turn is being killed so a checkpoint barrier can be admitted.
45pub(super) const RESTART_FOR_CHECKPOINT: WorkerRestartMessages = WorkerRestartMessages {
46    stop: "stop wedged Mjolnir worker before retrying checkpoint",
47    replace: "replace Mjolnir worker binary before retrying checkpoint",
48    start: "start Mjolnir worker after interrupting a wedged ACP turn",
49    connect: "connect to Mjolnir worker after restarting it for checkpoint",
50    project_memory: "project memory will not be synchronized after checkpoint worker restart",
51    native_session: "wait for ACP session after restarting the worker for checkpoint",
52};
53
54/// A quiet session is being moved onto the worker binary this controller
55/// would install.
56const RESTART_FOR_UPGRADE: WorkerRestartMessages = WorkerRestartMessages {
57    stop: "stop the Mjolnir worker before installing the current binary",
58    replace: "install the current Mjolnir worker binary",
59    start: "start Mjolnir worker on the current binary",
60    connect: "connect to Mjolnir worker after upgrading its binary",
61    project_memory: "project memory will not be synchronized after the worker upgrade",
62    native_session: "wait for ACP session after upgrading the worker",
63};
64
65/// What an upgrade attempt found. Nothing here is a failure: a worker that is
66/// already current and a session that started working again are both ordinary.
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub enum WorkerUpgradeOutcome {
69    /// The worker was replaced and the managed session now speaks to one
70    /// running this build.
71    Upgraded { build: String },
72    /// The worker already runs the binary this controller would install.
73    AlreadyCurrent { build: String },
74    /// The session was working when the upgrade reached it. A worker restart
75    /// would have killed that work, so nothing was touched.
76    Deferred,
77}
78
79impl WorkerUpgradeOutcome {
80    /// The build the session's worker runs now, or `None` when the attempt
81    /// stood down without establishing one.
82    #[must_use]
83    pub fn build(&self) -> Option<&str> {
84        match self {
85            Self::Upgraded { build } | Self::AlreadyCurrent { build } => Some(build),
86            Self::Deferred => None,
87        }
88    }
89}
90
91/// Whether a worker that reported `reported` in hello is running `installed`,
92/// the binary this controller would provision.
93///
94/// A worker that reported nothing is not: the field postdates it, so its
95/// binary does too.
96fn worker_runs_installed_build(reported: Option<&str>, installed: &str) -> bool {
97    reported.is_some_and(|reported| reported == installed)
98}
99
100impl Controller {
101    /// Replace a session's worker with the binary this controller would
102    /// install, when the session is quiet and its worker is a different build.
103    ///
104    /// `reported_build` is the digest the worker gave the observer that asked
105    /// for this. It only saves work: a match returns before anything is leased.
106    /// The decision that matters is taken again under the lease, against a
107    /// snapshot read from the worker itself, because a session can start
108    /// working between an observation and this call.
109    pub async fn upgrade_session_worker(
110        &self,
111        session_id: &str,
112        executor: &(impl CommandExecutor + Sync),
113        manager: &SessionManagerControl,
114        reported_build: Option<&str>,
115    ) -> Result<WorkerUpgradeOutcome> {
116        let (backend, worker_root) = self.worker_placement(session_id)?;
117        let reconnect = hel_targets::reconnect_plan(&backend, session_id)?
118            .commands
119            .into_iter()
120            .next()
121            .context("reconnect plan is empty")?;
122        let binary = worker_binary_for(&backend, executor)
123            .context("resolve the worker binary this controller would install")?;
124        let installed = hel::hel_worker_launch::worker_executable_digest(&binary)?;
125        if worker_runs_installed_build(reported_build, &installed) {
126            return Ok(WorkerUpgradeOutcome::AlreadyCurrent { build: installed });
127        }
128
129        // From here the session actor holds no connection, so no prompt it
130        // accepts can reach the worker: submissions queue until the lease
131        // returns. That is what makes the quiet check below decisive.
132        let handle = manager
133            .wait_for_session(session_id, UPGRADE_LEASE_TIMEOUT)
134            .await?;
135        let mut lease = handle.lease_connection().await?;
136        let snapshot = lease
137            .connection_mut()
138            .sync()
139            .await
140            .context("read the session state before upgrading its worker")?;
141        if worker_runs_installed_build(snapshot.worker_build.as_deref(), &installed) {
142            lease.release();
143            return Ok(WorkerUpgradeOutcome::AlreadyCurrent { build: installed });
144        }
145        if !snapshot.operational.is_quiet() {
146            lease.release();
147            return Ok(WorkerUpgradeOutcome::Deferred);
148        }
149
150        let restarted = self
151            .restart_worker_with_installed_binary(
152                session_id,
153                executor,
154                &backend,
155                &worker_root,
156                &reconnect,
157                &RESTART_FOR_UPGRADE,
158            )
159            .await;
160        match restarted {
161            Ok(connection) => {
162                lease.replace_connection(connection);
163                lease.release();
164                Ok(WorkerUpgradeOutcome::Upgraded { build: installed })
165            }
166            Err(error) => {
167                // Dropping the lease returns the actor to reconnecting on its
168                // own, which is the recovery for a half-finished restart.
169                drop(lease);
170                Err(error)
171            }
172        }
173    }
174
175    /// Stop the worker, install the binary this controller would provision,
176    /// start it and reconnect to the session it recovers.
177    pub(super) async fn restart_worker_with_installed_binary(
178        &self,
179        session_id: &str,
180        executor: &(impl CommandExecutor + Sync),
181        backend: &hel_targets::TargetLocator,
182        worker_root: &str,
183        reconnect: &CommandSpec,
184        messages: &WorkerRestartMessages,
185    ) -> Result<StandaloneSession> {
186        stop_worker_after_target_recovery(executor, backend, session_id, worker_root)
187            .context(messages.stop)?;
188        // Copy through hel.next and rename. scp/cp onto a still-mapped hel
189        // fails with ETXTBSY ("dest open ... Failure") even after SIGKILL,
190        // and prepare_worker_files writes that path in place.
191        let binary = worker_binary_for(backend, executor)?;
192        replace_installed_worker_binary(executor, backend, session_id, &binary)
193            .context(messages.replace)?;
194        start_worker(executor, backend, worker_root).context(messages.start)?;
195        // Journal recovery runs before the daemon binds control.sock. A long
196        // kimi session can take well over the ordinary 30s startup window.
197        let mut connection = match connect_started_worker_with_timeout(
198            reconnect,
199            session_id,
200            executor,
201            backend,
202            worker_root,
203            WORKER_RESTART_TIMEOUT,
204        )
205        .await
206        {
207            Ok(connection) => connection,
208            Err(error) => {
209                return Err(
210                    worker_probe_diagnosis(executor, backend, worker_root, error)
211                        .context(messages.connect),
212                );
213            }
214        };
215        let project_memory = match self.project_memory_sync_target(session_id) {
216            Ok(target) => Some(target),
217            Err(error) => {
218                tracing::warn!(
219                    session_id,
220                    error = format!("{error:#}"),
221                    "{}",
222                    messages.project_memory
223                );
224                None
225            }
226        };
227        connection.set_project_memory_target(project_memory);
228        wait_for_native_session(&mut connection, executor)
229            .await
230            .context(messages.native_session)?;
231        wait_for_idle_projection(&mut connection, WORKER_RESTART_TIMEOUT)
232            .await
233            .context("wait for ACP to go idle after worker restart")?;
234        Ok(connection)
235    }
236}
237
238/// Wait until a restarted worker's projection stops moving and reports idle.
239///
240/// Three stable polls, not one: a worker that has just recovered its journal
241/// can report idle between two events it is still applying.
242async fn wait_for_idle_projection(relay: &mut StandaloneSession, timeout: Duration) -> Result<()> {
243    let deadline = tokio::time::Instant::now() + timeout;
244    let mut last_ordinal = None;
245    let mut stable_polls = 0_u8;
246    loop {
247        let snapshot = relay.sync().await?;
248        let ordinal = snapshot.operational.latest_ordinal;
249        let idle = snapshot.operational.execution == RelayExecutionState::Idle;
250        if idle && last_ordinal == Some(ordinal) {
251            stable_polls = stable_polls.saturating_add(1);
252            if stable_polls >= 3 {
253                return Ok(());
254            }
255        } else {
256            stable_polls = 0;
257        }
258        last_ordinal = Some(ordinal);
259        if snapshot.operational.execution == RelayExecutionState::Closed {
260            bail!("ACP runtime stopped before becoming idle");
261        }
262        if tokio::time::Instant::now() >= deadline {
263            bail!(
264                "ACP runtime did not become idle after worker restart (execution={:?}, ordinal={ordinal})",
265                snapshot.operational.execution
266            );
267        }
268        tokio::time::sleep(Duration::from_millis(200)).await;
269    }
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    /// The three answers hello can produce, and what each means for the
277    /// worker's binary.
278    #[test]
279    fn only_a_matching_reported_build_counts_as_current() {
280        let installed = "a".repeat(64);
281
282        assert!(worker_runs_installed_build(Some(&installed), &installed));
283        assert!(!worker_runs_installed_build(
284            Some(&"b".repeat(64)),
285            &installed
286        ));
287        assert!(
288            !worker_runs_installed_build(None, &installed),
289            "a worker too old to report a build is older than this controller"
290        );
291    }
292}
mj_controller/hel_controller/worker_restart.rs

mj_controller/hel_controller/
worker_restart.rs
