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

mj_controller/hel_controller/
worker_restart.rs
