mj_controller/controller/checkpoint/
workspace_lease.rs

1use super::*;
2
3/// An idle workspace operation holds the managed connection and a worker
4/// barrier. Dropping this value disconnects and cancels the barrier; releasing
5/// it resumes dispatch without claiming that an archive covers the journal.
6pub struct IdleWorkspaceLease {
7    pub(super) lease: ManagedSessionLease,
8    pub(super) command_id: String,
9    pub(super) harness: HarnessKind,
10}
11
12impl IdleWorkspaceLease {
13    pub async fn acquire(handle: &ManagedSessionHandle, harness: HarnessKind) -> Result<Self> {
14        tokio::time::timeout(Duration::from_secs(30), async {
15            let mut lease = handle.lease_connection().await?;
16            let snapshot = lease.connection_mut().sync().await?;
17            ensure!(
18                snapshot.operational.safe_to_replace(harness),
19                "session must be live and idle with no queued or background work"
20            );
21            let command_id = new_command_id("workspace-write")?;
22            lease
23                .connection_mut()
24                .submit(
25                    command_id.clone(),
26                    RelayCommand::BeginCheckpoint {
27                        reason: Some("API workspace file write".into()),
28                    },
29                )
30                .await?;
31            loop {
32                let snapshot = lease.connection_mut().sync().await?;
33                if checkpoint_barrier_is_ready(&snapshot, &command_id) {
34                    let mut operation = Self {
35                        lease,
36                        command_id,
37                        harness,
38                    };
39                    operation.verify().await?;
40                    return Ok(operation);
41                }
42                ensure!(
43                    snapshot.operational.execution != RelayExecutionState::Running,
44                    "session started work before the file barrier was ready"
45                );
46                tokio::time::sleep(Duration::from_millis(25)).await;
47            }
48        })
49        .await
50        .context("session did not become available for a file write within 30 seconds")?
51    }
52
53    pub async fn verify(&mut self) -> Result<()> {
54        let mut snapshot = self.lease.connection_mut().sync().await?;
55        ensure!(
56            checkpoint_barrier_is_ready(&snapshot, &self.command_id),
57            "file write lost its workspace barrier"
58        );
59        snapshot.operational.checkpoint_barrier = None;
60        // Commands may queue behind this barrier, but cannot begin until it
61        // releases. Their arrival does not invalidate an in-progress write.
62        snapshot.operational.queued_prompts.clear();
63        ensure!(
64            snapshot.operational.safe_to_replace(self.harness),
65            "session is no longer idle for the file write"
66        );
67        Ok(())
68    }
69
70    pub async fn release(mut self) -> Result<()> {
71        tokio::time::timeout(Duration::from_secs(30), async {
72            self.lease
73                .connection_mut()
74                .submit(
75                    new_command_id("workspace-release")?,
76                    RelayCommand::ReleaseCheckpoint {
77                        barrier_command_id: self.command_id.clone(),
78                    },
79                )
80                .await?;
81            loop {
82                let snapshot = self.lease.connection_mut().sync().await?;
83                if snapshot.operational.checkpoint_barrier.as_deref() != Some(&self.command_id) {
84                    return Ok::<_, anyhow::Error>(());
85                }
86                tokio::time::sleep(Duration::from_millis(25)).await;
87            }
88        })
89        .await
90        .context("release file write barrier timed out")??;
91        self.lease.release();
92        Ok(())
93    }
94}
95
96pub(super) fn checkpoint_barrier_is_ready(
97    snapshot: &ManagedSessionSnapshot,
98    command_id: &str,
99) -> bool {
100    snapshot.operational.checkpoint_barrier.as_deref() == Some(command_id)
101        && snapshot.operational.checkpoint_ready.is_some()
102}
mj_controller/controller/checkpoint/workspace_lease.rs

mj_controller/controller/checkpoint/
workspace_lease.rs
