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 has a different compatibility
7//! identity from the 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    /// Worker-input compatibility identity 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            no_project_memory: false,
371            no_mailbox: false,
372            target_template_id: "podman".into(),
373            resource_allocation: None,
374            additional_mounts: Vec::new(),
375            state,
376            target: None,
377            native_session_id: None,
378            acp_session_title: None,
379            session_title_override: None,
380            created_at: "2026-08-09T12:00:00Z".into(),
381            updated_at: "2026-08-09T12:01:00Z".into(),
382            viewed_through_event_ordinal: 0,
383            draft_input: String::new(),
384            last_error: None,
385            last_checkpoint_error: None,
386            checkpoint: None,
387        }
388    }
389
390    fn observation(worker_build: Option<&str>, quiet: bool) -> WorkerUpgradeObservation {
391        WorkerUpgradeObservation {
392            session: session_record(SessionState::Running),
393            config: Config::default(),
394            worker_build: worker_build.map(str::to_owned),
395            quiet,
396        }
397    }
398
399    fn failure(detail: &str) -> WorkerUpgradeResult {
400        WorkerUpgradeResult {
401            session_id: "session-1".into(),
402            outcome: Err(detail.into()),
403            preparation_failure: None,
404            observed_updated_at: "2026-08-09T12:01:00Z".into(),
405            cancelled: false,
406        }
407    }
408
409    fn current(build: &str) -> WorkerUpgradeOutcome {
410        WorkerUpgradeOutcome::AlreadyCurrent {
411            build: build.to_owned(),
412        }
413    }
414
415    fn upgraded(build: &str) -> WorkerUpgradeOutcome {
416        WorkerUpgradeOutcome::Upgraded {
417            build: build.to_owned(),
418        }
419    }
420
421    fn success(outcome: WorkerUpgradeOutcome) -> WorkerUpgradeResult {
422        WorkerUpgradeResult {
423            session_id: "session-1".into(),
424            outcome: Ok(outcome),
425            preparation_failure: None,
426            observed_updated_at: "2026-08-09T12:01:00Z".into(),
427            cancelled: false,
428        }
429    }
430
431    #[test]
432    fn a_busy_turn_is_never_upgraded_no_matter_how_long_it_runs() {
433        let started = Utc::now();
434        let policy = PolicyState::default();
435        let two_days_later = started + chrono::Duration::days(2);
436
437        assert!(!policy.due(&observation(Some("old-build"), false), two_days_later));
438        assert!(
439            policy.due(&observation(Some("old-build"), true), two_days_later),
440            "the next quiet observation may upgrade without an age-based busy timeout"
441        );
442        assert!(
443            policy.due(&observation(None, true), two_days_later),
444            "a worker too old to report a build is outdated"
445        );
446    }
447
448    /// A session that is closing, checkpointing or already stopped is not a
449    /// session whose worker may be replaced underneath it.
450    #[test]
451    fn only_a_running_session_is_due() {
452        let now = Utc::now();
453        let policy = PolicyState::default();
454        for state in [
455            SessionState::Provisioning,
456            SessionState::Disconnected,
457            SessionState::Checkpointing,
458            SessionState::Closing,
459            SessionState::Destroying,
460            SessionState::Stopped,
461            SessionState::Lost,
462            SessionState::Error,
463            SessionState::DestroyedWithDataLoss,
464        ] {
465            let mut observation = observation(Some("build-a"), true);
466            observation.session.state = state;
467            assert!(!policy.due(&observation, now), "{state:?}");
468        }
469    }
470
471    /// A proved build stays current until a different worker is observed.
472    #[test]
473    fn a_worker_proved_current_stays_trusted_for_coordinator_lifetime() {
474        let now = Utc::now();
475        let mut policy = PolicyState::default();
476        policy.record(&success(current("build-a")), now);
477
478        assert!(!policy.due(&observation(Some("build-a"), true), now));
479        assert!(
480            !policy.due(
481                &observation(Some("build-a"), true),
482                now + chrono::Duration::days(2)
483            ),
484            "the launched build remains trusted for the coordinator lifetime"
485        );
486        assert!(
487            policy.due(&observation(Some("build-b"), true), now),
488            "a different build is outdated however recently the last one was checked"
489        );
490    }
491
492    /// A failed upgrade waits, and waits longer each time, so a broken target
493    /// is not restarted on every sync tick.
494    #[test]
495    fn a_failed_upgrade_backs_off_and_widens() {
496        let now = Utc::now();
497        let mut policy = PolicyState::default();
498        policy.record(&failure("install the current Mjolnir worker binary"), now);
499
500        let interval = chrono::Duration::from_std(WORKER_UPGRADE_RETRY_INTERVAL).unwrap();
501        assert!(!policy.due(&observation(Some("build-a"), true), now));
502        assert!(!policy.due(
503            &observation(Some("build-a"), true),
504            now + interval - chrono::Duration::seconds(1)
505        ));
506        assert!(policy.due(&observation(Some("build-a"), true), now + interval));
507
508        policy.record(&failure("install the current Mjolnir worker binary"), now);
509        assert!(!policy.due(
510            &observation(Some("build-a"), true),
511            now + interval * 2 - chrono::Duration::seconds(1)
512        ));
513        assert!(policy.due(&observation(Some("build-a"), true), now + interval * 2));
514    }
515
516    /// A successful upgrade stops the session being probed again: the worker
517    /// that answers next is the new one, and confirming it clears the failure
518    /// run the upgrade itself was guarded by.
519    #[test]
520    fn a_successful_upgrade_stops_the_probing_and_confirming_it_clears_the_backoff() {
521        let now = Utc::now();
522        let mut policy = PolicyState::default();
523        policy.record(&failure("install the current Mjolnir worker binary"), now);
524        policy.record(&success(upgraded("build-b")), now);
525
526        let confirmed = observation(Some("build-b"), true);
527        policy.observe(&confirmed);
528        assert_eq!(policy.consecutive_failures, 0);
529        assert_eq!(policy.failed_at, None);
530        assert!(
531            !policy.due(&confirmed, now),
532            "the worker now runs the installed build, so nothing is due"
533        );
534    }
535
536    /// An upgrade that does not take - the worker comes back reporting a build
537    /// that is still not the installed one - must not restart that worker on
538    /// every freshness window forever.
539    #[test]
540    fn an_upgrade_that_does_not_take_backs_off_instead_of_looping() {
541        let now = Utc::now();
542        let mut policy = PolicyState::default();
543        policy.record(&success(upgraded("build-b")), now);
544
545        // The worker came back as something else, so nothing confirms the
546        // upgrade and the cooldown stands.
547        let unchanged = observation(Some("build-a"), true);
548        policy.observe(&unchanged);
549        let interval = chrono::Duration::from_std(WORKER_UPGRADE_RETRY_INTERVAL).unwrap();
550        assert!(!policy.due(&unchanged, now));
551        assert!(policy.due(&unchanged, now + interval));
552
553        policy.record(&success(upgraded("build-b")), now + interval);
554        policy.observe(&unchanged);
555        assert!(!policy.due(&unchanged, now + interval * 2));
556        assert!(policy.due(&unchanged, now + interval * 3));
557    }
558
559    /// Deferrals retain a bounded retry deadline even if observations stay quiet.
560    #[test]
561    fn a_deferred_upgrade_waits_before_repeating_preparation() {
562        let now = Utc::now();
563        let mut policy = PolicyState::default();
564        policy.record(&success(WorkerUpgradeOutcome::Deferred), now);
565
566        assert!(!policy.due(&observation(Some("build-a"), true), now));
567        let delay = chrono::Duration::from_std(WORKER_UPGRADE_DEFERRED_INTERVAL).unwrap();
568        assert!(!policy.due(
569            &observation(Some("build-a"), true),
570            now + delay - chrono::Duration::milliseconds(1)
571        ));
572        assert!(policy.due(&observation(Some("build-a"), true), now + delay));
573    }
574
575    /// A preempted attempt says nothing either way: it must not count as a
576    /// failure and must not delay the next attempt.
577    #[test]
578    fn a_preempted_attempt_is_neither_a_success_nor_a_failure() {
579        let now = Utc::now();
580        let mut policy = PolicyState::default();
581        policy.record(
582            &WorkerUpgradeResult {
583                session_id: "session-1".into(),
584                outcome: Err("operation cancelled".into()),
585                preparation_failure: None,
586                observed_updated_at: "2026-08-09T12:01:00Z".into(),
587                cancelled: true,
588            },
589            now,
590        );
591
592        assert_eq!(policy.consecutive_failures, 0);
593        assert!(policy.due(&observation(Some("build-a"), true), now));
594    }
595
596    /// Recovery holds the same per-session slot, so an upgrade cannot start
597    /// while a recovery copy is running for that session.
598    #[test]
599    fn the_shared_gate_keeps_an_upgrade_and_a_recovery_copy_apart() {
600        let gate = Arc::new(RecoveryGate::default());
601        let recovery_copy = gate.try_start("session-1").expect("the slot starts free");
602
603        assert!(gate.try_start("session-1").is_none());
604
605        drop(recovery_copy);
606        assert!(gate.try_start("session-1").is_some());
607    }
608
609    /// Observing must hand off and return: the daemon reports from its event
610    /// loop, and no upgrade decision may hold that loop up.
611    #[test]
612    fn observing_hands_off_without_waiting() {
613        let (observations, mut queued) = observation_channel();
614        let observer = WorkerUpgradeObserver { observations };
615
616        for _ in 0..64 {
617            observer.observe(observation(Some("build-a"), true));
618        }
619
620        let received = std::iter::from_fn(|| queued.try_recv()).count();
621        assert_eq!(received, 1);
622    }
623}
mj_controller/worker_upgrade.rs

mj_controller/
worker_upgrade.rs
