mj_controller/
worker_upgrade.rs

1//! Background worker-binary upgrade policy and coordination.
2//!
3//! A running session keeps the worker it started with, so a session that is
4//! never stopped never gains anything a newer daemon's worker learned. This
5//! coordinator watches the same session views the recovery coordinator does
6//! and, when a session is quiet and its worker is a different build from the
7//! one this controller would install, replaces it in place.
8//!
9//! Quiet is the whole safety argument: stopping a worker tears down the ACP
10//! bridge with it, so an upgrade may only run when nothing would be lost.
11
12use std::collections::BTreeMap;
13use std::sync::Arc;
14use std::sync::atomic::Ordering;
15use std::time::Duration;
16
17use chrono::{DateTime, Utc};
18use tokio::task::JoinSet;
19
20use crate::controller::{Controller, HarnessPreparationFailure, WorkerUpgradeOutcome};
21use crate::recovery::{backoff_delay, elapsed_at_least};
22use crate::recovery_gate::{
23    ObservationReceiver, ObservationSender, RecoveryGate, RecoveryObserver, observation_channel,
24};
25use crate::session_manager::SessionManagerControl;
26use crate::targets::CancellableProcessExecutor;
27use mj_core::config::Config;
28use mj_core::state::{SessionRecord, SessionState, State};
29
30/// How long a failed upgrade waits before it is tried again, doubling per
31/// consecutive failure. The daemon periodically reobserves quiet sessions, so
32/// without this one broken target would be retried continuously.
33const WORKER_UPGRADE_RETRY_INTERVAL: Duration = Duration::from_secs(10 * 60);
34const WORKER_UPGRADE_DEFERRED_INTERVAL: Duration = Duration::from_secs(30);
35
36/// Ceiling on the widening retry delay, so a target that is broken rather than
37/// blipping is still probed, just rarely.
38const MAX_WORKER_UPGRADE_RETRY_INTERVAL: Duration = Duration::from_secs(2 * 60 * 60);
39
40/// Upper bound on one upgrade. Stopping, installing, starting and waiting for
41/// a recovered journal all happen inside it; past this the attempt is a
42/// reported failure rather than a session whose upgrade never ends.
43const WORKER_UPGRADE_TIMEOUT: Duration = Duration::from_secs(15 * 60);
44
45/// One session view, reduced to what the upgrade decision reads.
46#[derive(Debug, Clone)]
47pub struct WorkerUpgradeObservation {
48    pub session: SessionRecord,
49    pub config: Config,
50    /// Content address the connected worker reported in hello, or `None` when
51    /// it is too old to report one. `None` counts as outdated.
52    pub worker_build: Option<String>,
53    /// Whether replacing the worker now would destroy nothing. See
54    /// [`mj_core::relay::RelayOperationalState::is_quiet`].
55    ///
56    /// An observer may skip reporting a session that is working - the daemon
57    /// does, to keep the config clone off the streaming path - but the rule
58    /// lives here, so a busy observation that does arrive is still refused.
59    pub quiet: bool,
60}
61
62/// Reports session activity to the upgrade coordinator.
63///
64/// Like the recovery observer, this is a queued hand-off: the caller is an
65/// event loop and must never wait on an upgrade decision.
66#[derive(Clone)]
67pub struct WorkerUpgradeObserver {
68    observations: ObservationSender<WorkerUpgradeObservation>,
69}
70
71impl WorkerUpgradeObserver {
72    pub fn observe(&self, observation: WorkerUpgradeObservation) {
73        self.observations
74            .send(observation.session.id.clone(), observation, |_, _| {});
75    }
76}
77
78#[derive(Debug, Clone)]
79pub struct WorkerUpgradeResult {
80    pub session_id: String,
81    pub outcome: Result<WorkerUpgradeOutcome, String>,
82    /// A failed harness preparation is a session failure, not a transient
83    /// worker-upgrade error. The daemon uses the observed revision to avoid
84    /// overwriting a newer lifecycle transition.
85    pub(crate) preparation_failure: Option<HarnessPreparationFailure>,
86    pub(crate) observed_updated_at: String,
87    /// The attempt was preempted by a lifecycle operation or by coordinator
88    /// shutdown. It judged nothing, so it is neither a success nor a failure.
89    pub cancelled: bool,
90}
91
92pub struct WorkerUpgradeCoordinator {
93    supervisor: Option<tokio::task::JoinHandle<()>>,
94    observer: WorkerUpgradeObserver,
95    results: ObservationReceiver<WorkerUpgradeResult>,
96    gate: Arc<RecoveryGate>,
97}
98
99impl Drop for WorkerUpgradeCoordinator {
100    fn drop(&mut self) {
101        // Closing the shared gate prevents new attempts and cancels existing
102        // ones atomically. Both supervisors retain ownership until settlement.
103        self.gate.close();
104    }
105}
106
107impl WorkerUpgradeCoordinator {
108    /// Start the coordinator, sharing the recovery observer's gate so a
109    /// recovery copy and an upgrade never touch one session at the same time.
110    pub fn spawn(session_manager: SessionManagerControl, recovery: &RecoveryObserver) -> Self {
111        let (observations_tx, mut observations_rx) =
112            observation_channel::<WorkerUpgradeObservation>();
113        let (results_tx, results_rx) = observation_channel();
114        let gate = recovery.gate.clone();
115        let coordinator_gate = gate.clone();
116        let supervisor = tokio::spawn(async move {
117            let mut policies = BTreeMap::<String, PolicyState>::new();
118            let mut attempts = JoinSet::new();
119            let mut closing = false;
120            loop {
121                if closing && attempts.is_empty() {
122                    break;
123                }
124                tokio::select! {
125                    _ = coordinator_gate.closed(), if !closing => { closing = true; }
126                    observed = observations_rx.recv(), if !closing => {
127                        let Some(observation) = observed else { coordinator_gate.close(); closing = true; continue; };
128                        let session_id = observation.session.id.clone();
129                        let policy = policies.entry(session_id.clone()).or_default();
130                        policy.observe(&observation);
131                        if !policy.due(&observation, Utc::now()) {
132                            continue;
133                        }
134                        let Some(upgrade_cancelled) = coordinator_gate.try_start(&session_id)
135                        else {
136                            continue;
137                        };
138                        let session_manager = session_manager.clone();
139                        let task_cancelled = upgrade_cancelled.cancellation();
140                        let task_session_id = session_id.clone();
141                        attempts.spawn(async move {
142                            let observed_updated_at = observation.session.updated_at.clone();
143                            let (joined, admission) = upgrade_cancelled.run_blocking(move |cancelled| {
144                                let Some(session) = crate::recovery_gate::current_background_session(&observation.session)
145                                    .map_err(|error| format!("read current upgrade placement: {error:#}"))?
146                                else { return Ok((Ok(WorkerUpgradeOutcome::Deferred), None)); };
147                                let mut state = State::default();
148                                state.sessions.insert(task_session_id.clone(), session);
149                                let controller = Controller {
150                                    config: observation.config,
151                                    state,
152                                };
153                                let executor = CancellableProcessExecutor::new(cancelled)
154                                    .with_deadline(WORKER_UPGRADE_TIMEOUT);
155                                let outcome = mj_core::runtime::block_on(controller.upgrade_session_worker(
156                                    &task_session_id,
157                                    &executor,
158                                    &session_manager,
159                                    observation.worker_build.as_deref(),
160                                ))
161                                .and_then(|result| result);
162                                let preparation_failure = outcome
163                                    .as_ref()
164                                    .err()
165                                    .and_then(|error| error.downcast_ref::<HarnessPreparationFailure>())
166                                    .cloned();
167                                Ok((
168                                    outcome.map_err(|error| format!("{error:#}")),
169                                    preparation_failure,
170                                ))
171                            })
172                            .await;
173                            let (outcome, preparation_failure) = match joined {
174                                Ok(Ok(result)) => result,
175                                Ok(Err(error)) => (Err(error), None),
176                                Err(error) => (
177                                    Err(format!("worker upgrade task failed: {error}")),
178                                    None,
179                                ),
180                            };
181                            let result = WorkerUpgradeResult {
182                                session_id,
183                                outcome,
184                                preparation_failure,
185                                observed_updated_at,
186                                cancelled: task_cancelled.load(Ordering::Acquire),
187                            };
188                            (result, admission)
189                        });
190                    }
191                    completed = attempts.join_next(), if !attempts.is_empty() => {
192                        let Some(completed) = completed else { continue };
193                        let (result, _admission) = match completed {
194                            Ok(completed) => completed,
195                            Err(error) => { tracing::error!(%error, "worker upgrade attempt supervisor failed"); continue; }
196                        };
197                        let policy = policies.entry(result.session_id.clone()).or_default();
198                        policy.record(&result, Utc::now());
199                        if let Err(error) = &result.outcome && !result.cancelled {
200                            tracing::warn!(session_id = %result.session_id, %error, "background worker upgrade failed");
201                        }
202                        results_tx.send(result.session_id.clone(), result, |_, _| {});
203                    }
204                }
205            }
206        });
207        Self {
208            supervisor: Some(supervisor),
209            observer: WorkerUpgradeObserver {
210                observations: observations_tx,
211            },
212            results: results_rx,
213            gate,
214        }
215    }
216
217    /// Close admission and wait for every admitted executor and durable
218    /// settlement. The daemon applies its shared shutdown watchdog outside.
219    pub async fn shutdown(&mut self) -> anyhow::Result<()> {
220        self.gate.close();
221        if let Some(supervisor) = self.supervisor.take() {
222            supervisor.await.map_err(|error| {
223                anyhow::anyhow!("background coordinator supervisor failed: {error}")
224            })?;
225        }
226        Ok(())
227    }
228
229    pub fn observer(&self) -> WorkerUpgradeObserver {
230        self.observer.clone()
231    }
232
233    pub fn try_result(&mut self) -> Option<WorkerUpgradeResult> {
234        self.results.try_recv()
235    }
236}
237
238/// What the coordinator remembers about one session between observations.
239#[derive(Debug, Default, PartialEq, Eq)]
240struct PolicyState {
241    /// The build proved current for this coordinator. While the observed
242    /// build still matches it, no attempt is needed and nothing is hashed.
243    current_build: Option<String>,
244    failed_at: Option<DateTime<Utc>>,
245    deferred_at: Option<DateTime<Utc>>,
246    consecutive_failures: u32,
247}
248
249impl PolicyState {
250    /// Whether an upgrade attempt should start for this observation.
251    fn due(&self, observation: &WorkerUpgradeObservation, now: DateTime<Utc>) -> bool {
252        // Killing a worker mid-turn destroys the turn, and a session that is
253        // shutting down or already stopped has no worker worth replacing.
254        if !observation.quiet || observation.session.state != SessionState::Running {
255            return false;
256        }
257        if self.worker_is_known_current(observation.worker_build.as_deref()) {
258            return false;
259        }
260        if self.deferred_at.is_some_and(|deferred| {
261            !elapsed_at_least(deferred, now, WORKER_UPGRADE_DEFERRED_INTERVAL)
262        }) {
263            return false;
264        }
265        self.failed_at.is_none_or(|failed_at| {
266            elapsed_at_least(
267                failed_at,
268                now,
269                backoff_delay(
270                    WORKER_UPGRADE_RETRY_INTERVAL,
271                    MAX_WORKER_UPGRADE_RETRY_INTERVAL,
272                    self.consecutive_failures,
273                ),
274            )
275        })
276    }
277
278    /// Whether a previous attempt proved this exact build current. A worker
279    /// reporting no build is never current: it predates the field, so it
280    /// predates this controller.
281    fn worker_is_known_current(&self, worker_build: Option<&str>) -> bool {
282        let (Some(observed), Some(current)) = (worker_build, self.current_build.as_deref()) else {
283            return false;
284        };
285        observed == current
286    }
287
288    /// Fold in what one observation proves, before deciding whether to act.
289    ///
290    /// The one thing it can prove is that the last upgrade took: the worker
291    /// now reports the build that upgrade installed. That releases the
292    /// cooldown an upgrade leaves behind.
293    fn observe(&mut self, observation: &WorkerUpgradeObservation) {
294        if self.current_build.is_some()
295            && observation.worker_build.as_deref() == self.current_build.as_deref()
296        {
297            self.failed_at = None;
298            self.deferred_at = None;
299            self.consecutive_failures = 0;
300        }
301    }
302
303    fn record(&mut self, result: &WorkerUpgradeResult, now: DateTime<Utc>) {
304        if result.cancelled {
305            // A preempted attempt judged nothing: it must neither be counted
306            // as a failure nor suppress the next observation.
307            return;
308        }
309        self.deferred_at = None;
310        match &result.outcome {
311            Ok(WorkerUpgradeOutcome::Deferred) => {
312                // A busy worker or target recovery can refuse the swap. Do not
313                // repeat preparation on every unchanged quiet observation.
314                self.deferred_at = Some(now);
315            }
316            Ok(outcome @ WorkerUpgradeOutcome::AlreadyCurrent { .. }) => {
317                self.failed_at = None;
318                self.consecutive_failures = 0;
319                self.current_build = outcome.build().map(str::to_owned);
320            }
321            Ok(outcome @ WorkerUpgradeOutcome::Upgraded { .. }) => {
322                // The worker this session runs now, so the next observation
323                // reporting it needs no attempt - and, through `observe`,
324                // clears the cooldown below.
325                self.current_build = outcome.build().map(str::to_owned);
326                // An upgrade that did not take would otherwise restart this
327                // worker on every sync tick, forever. The same widening
328                // cooldown a failure gets bounds that; a worker that comes
329                // back reporting the installed build releases it immediately,
330                // so a healthy upgrade pays nothing.
331                self.consecutive_failures = self.consecutive_failures.saturating_add(1);
332                self.failed_at = Some(now);
333            }
334            Err(_) => {
335                self.consecutive_failures = self.consecutive_failures.saturating_add(1);
336                self.failed_at = Some(now);
337            }
338        }
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345
346    fn session_record(state: SessionState) -> SessionRecord {
347        SessionRecord {
348            project: None,
349            target_runtime: None,
350            launch_base: None,
351            launch_branch: None,
352            checkout: None,
353            publication: None,
354            build_cache: None,
355            container_workspace: None,
356            subagents: None,
357            create_managed_worktree: None,
358            workspace_id: mj_core::workspace::DEFAULT_WORKSPACE_ID.to_owned(),
359            archived: false,
360            container_cpus: None,
361            container_memory: None,
362            id: "session-1".to_owned(),
363            title: "work".into(),
364            harness_kind: mj_core::config::HarnessKind::Codex,
365            last_profile: "codex-1".into(),
366            bundle_id: "hel".into(),
367            project_directory: None,
368            managed_worktree: None,
369            review: None,
370            target_template_id: "podman".into(),
371            resource_allocation: None,
372            additional_mounts: Vec::new(),
373            state,
374            target: None,
375            native_session_id: None,
376            acp_session_title: None,
377            session_title_override: None,
378            created_at: "2026-08-09T12:00:00Z".into(),
379            updated_at: "2026-08-09T12:01:00Z".into(),
380            viewed_through_event_ordinal: 0,
381            draft_input: String::new(),
382            last_error: None,
383            last_checkpoint_error: None,
384            checkpoint: None,
385        }
386    }
387
388    fn observation(worker_build: Option<&str>, quiet: bool) -> WorkerUpgradeObservation {
389        WorkerUpgradeObservation {
390            session: session_record(SessionState::Running),
391            config: Config::default(),
392            worker_build: worker_build.map(str::to_owned),
393            quiet,
394        }
395    }
396
397    fn failure(detail: &str) -> WorkerUpgradeResult {
398        WorkerUpgradeResult {
399            session_id: "session-1".into(),
400            outcome: Err(detail.into()),
401            preparation_failure: None,
402            observed_updated_at: "2026-08-09T12:01:00Z".into(),
403            cancelled: false,
404        }
405    }
406
407    fn current(build: &str) -> WorkerUpgradeOutcome {
408        WorkerUpgradeOutcome::AlreadyCurrent {
409            build: build.to_owned(),
410        }
411    }
412
413    fn upgraded(build: &str) -> WorkerUpgradeOutcome {
414        WorkerUpgradeOutcome::Upgraded {
415            build: build.to_owned(),
416        }
417    }
418
419    fn success(outcome: WorkerUpgradeOutcome) -> WorkerUpgradeResult {
420        WorkerUpgradeResult {
421            session_id: "session-1".into(),
422            outcome: Ok(outcome),
423            preparation_failure: None,
424            observed_updated_at: "2026-08-09T12:01:00Z".into(),
425            cancelled: false,
426        }
427    }
428
429    #[test]
430    fn a_busy_turn_is_never_upgraded_no_matter_how_long_it_runs() {
431        let started = Utc::now();
432        let policy = PolicyState::default();
433        let two_days_later = started + chrono::Duration::days(2);
434
435        assert!(!policy.due(&observation(Some("old-build"), false), two_days_later));
436        assert!(
437            policy.due(&observation(Some("old-build"), true), two_days_later),
438            "the next quiet observation may upgrade without an age-based busy timeout"
439        );
440        assert!(
441            policy.due(&observation(None, true), two_days_later),
442            "a worker too old to report a build is outdated"
443        );
444    }
445
446    /// A session that is closing, checkpointing or already stopped is not a
447    /// session whose worker may be replaced underneath it.
448    #[test]
449    fn only_a_running_session_is_due() {
450        let now = Utc::now();
451        let policy = PolicyState::default();
452        for state in [
453            SessionState::Provisioning,
454            SessionState::Disconnected,
455            SessionState::Checkpointing,
456            SessionState::Closing,
457            SessionState::Destroying,
458            SessionState::Stopped,
459            SessionState::Lost,
460            SessionState::Error,
461            SessionState::DestroyedWithDataLoss,
462        ] {
463            let mut observation = observation(Some("build-a"), true);
464            observation.session.state = state;
465            assert!(!policy.due(&observation, now), "{state:?}");
466        }
467    }
468
469    /// A proved build stays current until a different worker is observed.
470    #[test]
471    fn a_worker_proved_current_stays_trusted_for_coordinator_lifetime() {
472        let now = Utc::now();
473        let mut policy = PolicyState::default();
474        policy.record(&success(current("build-a")), now);
475
476        assert!(!policy.due(&observation(Some("build-a"), true), now));
477        assert!(
478            !policy.due(
479                &observation(Some("build-a"), true),
480                now + chrono::Duration::days(2)
481            ),
482            "the launched build remains trusted for the coordinator lifetime"
483        );
484        assert!(
485            policy.due(&observation(Some("build-b"), true), now),
486            "a different build is outdated however recently the last one was checked"
487        );
488    }
489
490    /// A failed upgrade waits, and waits longer each time, so a broken target
491    /// is not restarted on every sync tick.
492    #[test]
493    fn a_failed_upgrade_backs_off_and_widens() {
494        let now = Utc::now();
495        let mut policy = PolicyState::default();
496        policy.record(&failure("install the current Mjolnir worker binary"), now);
497
498        let interval = chrono::Duration::from_std(WORKER_UPGRADE_RETRY_INTERVAL).unwrap();
499        assert!(!policy.due(&observation(Some("build-a"), true), now));
500        assert!(!policy.due(
501            &observation(Some("build-a"), true),
502            now + interval - chrono::Duration::seconds(1)
503        ));
504        assert!(policy.due(&observation(Some("build-a"), true), now + interval));
505
506        policy.record(&failure("install the current Mjolnir worker binary"), now);
507        assert!(!policy.due(
508            &observation(Some("build-a"), true),
509            now + interval * 2 - chrono::Duration::seconds(1)
510        ));
511        assert!(policy.due(&observation(Some("build-a"), true), now + interval * 2));
512    }
513
514    /// A successful upgrade stops the session being probed again: the worker
515    /// that answers next is the new one, and confirming it clears the failure
516    /// run the upgrade itself was guarded by.
517    #[test]
518    fn a_successful_upgrade_stops_the_probing_and_confirming_it_clears_the_backoff() {
519        let now = Utc::now();
520        let mut policy = PolicyState::default();
521        policy.record(&failure("install the current Mjolnir worker binary"), now);
522        policy.record(&success(upgraded("build-b")), now);
523
524        let confirmed = observation(Some("build-b"), true);
525        policy.observe(&confirmed);
526        assert_eq!(policy.consecutive_failures, 0);
527        assert_eq!(policy.failed_at, None);
528        assert!(
529            !policy.due(&confirmed, now),
530            "the worker now runs the installed build, so nothing is due"
531        );
532    }
533
534    /// An upgrade that does not take - the worker comes back reporting a build
535    /// that is still not the installed one - must not restart that worker on
536    /// every freshness window forever.
537    #[test]
538    fn an_upgrade_that_does_not_take_backs_off_instead_of_looping() {
539        let now = Utc::now();
540        let mut policy = PolicyState::default();
541        policy.record(&success(upgraded("build-b")), now);
542
543        // The worker came back as something else, so nothing confirms the
544        // upgrade and the cooldown stands.
545        let unchanged = observation(Some("build-a"), true);
546        policy.observe(&unchanged);
547        let interval = chrono::Duration::from_std(WORKER_UPGRADE_RETRY_INTERVAL).unwrap();
548        assert!(!policy.due(&unchanged, now));
549        assert!(policy.due(&unchanged, now + interval));
550
551        policy.record(&success(upgraded("build-b")), now + interval);
552        policy.observe(&unchanged);
553        assert!(!policy.due(&unchanged, now + interval * 2));
554        assert!(policy.due(&unchanged, now + interval * 3));
555    }
556
557    /// Deferrals retain a bounded retry deadline even if observations stay quiet.
558    #[test]
559    fn a_deferred_upgrade_waits_before_repeating_preparation() {
560        let now = Utc::now();
561        let mut policy = PolicyState::default();
562        policy.record(&success(WorkerUpgradeOutcome::Deferred), now);
563
564        assert!(!policy.due(&observation(Some("build-a"), true), now));
565        let delay = chrono::Duration::from_std(WORKER_UPGRADE_DEFERRED_INTERVAL).unwrap();
566        assert!(!policy.due(
567            &observation(Some("build-a"), true),
568            now + delay - chrono::Duration::milliseconds(1)
569        ));
570        assert!(policy.due(&observation(Some("build-a"), true), now + delay));
571    }
572
573    /// A preempted attempt says nothing either way: it must not count as a
574    /// failure and must not delay the next attempt.
575    #[test]
576    fn a_preempted_attempt_is_neither_a_success_nor_a_failure() {
577        let now = Utc::now();
578        let mut policy = PolicyState::default();
579        policy.record(
580            &WorkerUpgradeResult {
581                session_id: "session-1".into(),
582                outcome: Err("operation cancelled".into()),
583                preparation_failure: None,
584                observed_updated_at: "2026-08-09T12:01:00Z".into(),
585                cancelled: true,
586            },
587            now,
588        );
589
590        assert_eq!(policy.consecutive_failures, 0);
591        assert!(policy.due(&observation(Some("build-a"), true), now));
592    }
593
594    /// Recovery holds the same per-session slot, so an upgrade cannot start
595    /// while a recovery copy is running for that session.
596    #[test]
597    fn the_shared_gate_keeps_an_upgrade_and_a_recovery_copy_apart() {
598        let gate = Arc::new(RecoveryGate::default());
599        let recovery_copy = gate.try_start("session-1").expect("the slot starts free");
600
601        assert!(gate.try_start("session-1").is_none());
602
603        drop(recovery_copy);
604        assert!(gate.try_start("session-1").is_some());
605    }
606
607    /// Observing must hand off and return: the daemon reports from its event
608    /// loop, and no upgrade decision may hold that loop up.
609    #[test]
610    fn observing_hands_off_without_waiting() {
611        let (observations, mut queued) = observation_channel();
612        let observer = WorkerUpgradeObserver { observations };
613
614        for _ in 0..64 {
615            observer.observe(observation(Some("build-a"), true));
616        }
617
618        let received = std::iter::from_fn(|| queued.try_recv()).count();
619        assert_eq!(received, 1);
620    }
621}
mj_controller/worker_upgrade.rs

mj_controller/
worker_upgrade.rs
