run: Support labeled revisions.

CHANGELOG.md

@@ -33,6 +33,11 @@ to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
   back to prompting the user if the heuristics are inconclusive. It can also run
   in non-interactive mode, which aborts if prompting would be needed.
 
+* `jj run` now supports labeled revisions (e.g. `-r before=trunk() -r after=@`).
+  The label is exposed to the executed command via the `$JJ_LABEL` environment
+  variable, and allocates a dedicated workspace slot to preserve build cache
+  affinity.
+
 ### Fixed bugs
 
 * The default pager flags now include `-K` (`--quit-on-intr`), so pressing

cli/src/commands/run.rs

@@ -28,6 +28,7 @@ use std::process::ExitStatus;
 use std::process::Output;
 use std::process::Stdio;
 use std::sync::Arc;
+use std::sync::LazyLock;
 use std::time::Duration;
 
 use futures::StreamExt as _;
@@ -36,7 +37,6 @@ use itertools::Itertools as _;
 use jj_lib::backend::BackendError;
 use jj_lib::backend::CommitId;
 use jj_lib::commit::Commit;
-use jj_lib::commit::CommitIteratorExt as _;
 use jj_lib::conflicts::ConflictMarkerStyle;
 use jj_lib::fsmonitor::FsmonitorSettings;
 use jj_lib::gitignore::GitIgnoreFile;
@@ -55,6 +55,7 @@ use jj_lib::merged_tree::MergedTree;
 use jj_lib::object_id::ObjectId as _;
 use jj_lib::repo::Repo as _;
 use jj_lib::working_copy::SnapshotOptions;
+use regex::Regex;
 use tokio::runtime::Builder;
 use tokio::sync::mpsc;
 use tokio::sync::mpsc::Sender;
@@ -68,6 +69,7 @@ use crate::cli_util::WorkspaceCommandHelper;
 use crate::cli_util::WorkspaceCommandTransaction;
 use crate::command_error::CommandError;
 use crate::command_error::CommandErrorKind;
+use crate::command_error::user_error;
 use crate::ui::Ui;
 
 #[derive(Debug, thiserror::Error)]
@@ -190,21 +192,22 @@ impl WorkspacePool {
     async fn acquire(
         &self,
         commit: &Commit,
+        label: Option<&str>,
         base_ignores: Arc<GitIgnoreFile>,
     ) -> Result<RunWorkspace, RunError> {
         // Find a free slot. The first iteration may fail if another worker
         // (this process or another) holds every slot; sleep and retry.
         let mut cur_sleep = Duration::from_millis(10);
         let max_sleep = Duration::from_millis(250);
-        let (slot_index, lock) = loop {
-            if let Some(found) = self.try_acquire_any_slot()? {
+        let (slot_name, lock) = loop {
+            if let Some(found) = self.try_acquire_slot(label)? {
                 break found;
             }
             sleep(cur_sleep).await;
             cur_sleep = min(cur_sleep.saturating_mul(2), max_sleep);
         };
 
-        let slot_path = self.slot_path(slot_index);
+        let slot_path = self.slot_path(&slot_name);
         let working_copy_dir = slot_path.join("working_copy");
         let state_dir = slot_path.join("state");
         let tree_state_path = state_dir.join("tree_state");
@@ -325,25 +328,40 @@ impl WorkspacePool {
         }
     }
 
-    fn slot_path(&self, index: usize) -> PathBuf {
-        self.base_path.join(index.to_string())
+    fn slot_path(&self, slot_name: &str) -> PathBuf {
+        self.base_path.join(slot_name)
     }
 
-    fn slot_lock_path(&self, index: usize) -> PathBuf {
-        self.base_path.join(format!("{index}.lock"))
+    fn slot_lock_path(&self, slot_name: &str) -> PathBuf {
+        self.base_path.join(format!("{slot_name}.lock"))
     }
 
-    /// Try to acquire any slot's lock without blocking. Returns the slot
-    /// index and the held lock if one was available, `Ok(None)` if every
+    /// Try to acquire a slot's lock without blocking. Returns the slot
+    /// name and the held lock if one was available, `Ok(None)` if every
     /// slot was contended.
-    fn try_acquire_any_slot(&self) -> Result<Option<(usize, FileLock)>, RunError> {
+    fn try_acquire_slot(
+        &self,
+        label: Option<&str>,
+    ) -> Result<Option<(String, FileLock)>, RunError> {
+        if let Some(label) = label {
+            let slot_path = self.slot_path(label);
+            fs::create_dir_all(&slot_path)
+                .map_err(|e| RunError::PathCreationFailure(slot_path.clone(), e))?;
+            if let Some(lock) = FileLock::try_lock(self.slot_lock_path(label))? {
+                tracing::debug!(slot = label, "acquired labeled pool slot");
+                return Ok(Some((label.to_string(), lock)));
+            }
+            return Ok(None);
+        }
+
         for slot in 1..=self.size.get() {
-            let slot_path = self.slot_path(slot);
+            let slot_name = slot.to_string();
+            let slot_path = self.slot_path(&slot_name);
             fs::create_dir_all(&slot_path)
                 .map_err(|e| RunError::PathCreationFailure(slot_path.clone(), e))?;
-            if let Some(lock) = FileLock::try_lock(self.slot_lock_path(slot))? {
-                tracing::debug!(slot, "acquired pool slot");
-                return Ok(Some((slot, lock)));
+            if let Some(lock) = FileLock::try_lock(self.slot_lock_path(&slot_name))? {
+                tracing::debug!(slot = slot, "acquired pool slot");
+                return Ok(Some((slot_name, lock)));
             }
         }
         Ok(None)
@@ -382,7 +400,7 @@ async fn run_inner(
     handle: &tokio::runtime::Handle,
     spec: Arc<CommandSpec>,
     pool: Arc<WorkspacePool>,
-    commits: &[Commit],
+    commits: &[(Commit, Option<String>)],
     jobs: usize,
     passthrough: bool,
     ignore_errors: bool,
@@ -394,15 +412,18 @@ async fn run_inner(
         // Launch commits in order, keeping at most `jobs` in flight so tasks
         // start in commit order and the pool is never oversubscribed.
         while command_futures.len() < jobs {
-            let Some(commit) = commits_iter.next() else {
+            let Some((commit, label)) = commits_iter.next() else {
                 break;
             };
             let base_ignores = base_ignores.clone();
             let pool = pool.clone();
             let commit = commit.clone();
+            let label = label.clone();
             let spec = spec.clone();
             command_futures.spawn_on(
-                async move { rewrite_commit(base_ignores, pool, commit, spec, passthrough).await },
+                async move {
+                    rewrite_commit(base_ignores, pool, commit, label, spec, passthrough).await
+                },
                 handle,
             );
         }
@@ -440,10 +461,13 @@ async fn rewrite_commit(
     base_ignores: Arc<GitIgnoreFile>,
     pool: Arc<WorkspacePool>,
     commit: Commit,
+    label: Option<String>,
     spec: Arc<CommandSpec>,
     passthrough: bool,
 ) -> Result<RunJob, RunError> {
-    let mut workspace = pool.acquire(&commit, base_ignores.clone()).await?;
+    let mut workspace = pool
+        .acquire(&commit, label.as_deref(), base_ignores.clone())
+        .await?;
     let working_copy_dir = workspace.working_copy_dir.clone();
     let old_id = commit.id().clone();
     let old_tree = commit.tree();
@@ -496,6 +520,9 @@ async fn rewrite_commit(
         .env("JJ_COMMIT_ID", commit.id().hex())
         .stdin(Stdio::null())
         .kill_on_drop(true);
+    if let Some(ref l) = label {
+        command.env("JJ_LABEL", l);
+    }
     let output = if passthrough {
         // Connect stdout/stderr directly to the terminal so TTY-aware
         // programs work as expected. Capture is not possible; we wait
@@ -582,6 +609,7 @@ async fn rewrite_commit(
 /// - JJ_CHANGE_ID
 /// - JJ_COMMIT_ID
 /// - JJ_WORKSPACE_ROOT
+/// - JJ_LABEL (if a label was specified for the revision)
 ///
 /// ### Example
 ///
@@ -604,8 +632,23 @@ pub struct RunArgs {
     args: Vec<String>,
 
     /// The revisions to change
-    #[arg(long = "revision", short, value_name = "REVSETS", alias = "revisions")]
-    revisions: Vec<RevisionArg>,
+    ///
+    /// An optional label prefix can be specified, with at most one label per
+    /// revision (e.g. `-r before=trunk() -r after=@`). If it is specified:
+    ///
+    /// * It will be exposed to the command in `$JJ_LABEL`, allowing you to
+    ///   store outputs at known paths - for example: `jj run ... bash -c
+    ///   'run_build && cp out/binary /tmp/$JJ_LABEL'`
+    /// * It will run in a consistent workspace, ensuring cache affinity for
+    ///   incremental builds.
+    #[arg(
+        long = "revision",
+        short,
+        value_name = "REVSETS",
+        alias = "revisions",
+        verbatim_doc_comment
+    )]
+    revisions: Vec<String>,
 
     /// A no-op option to match the interface of `git rebase -x`
     #[arg(short = 'x', hide = true)]
@@ -697,6 +740,20 @@ fn resolve_jobs(
     Ok(NonZeroUsize::MIN)
 }
 
+static LABEL_RE: LazyLock<Regex> =
+    LazyLock::new(|| Regex::new(r"^([a-zA-Z_][a-zA-Z0-9_-]*)=(.*)$").unwrap());
+
+fn parse_labeled_revision(s: &str) -> (RevisionArg, Option<String>) {
+    if let Some(caps) = LABEL_RE.captures(s) {
+        (
+            RevisionArg::from(caps[2].to_string()),
+            Some(caps[1].to_string()),
+        )
+    } else {
+        (RevisionArg::from(s.to_string()), None)
+    }
+}
+
 pub async fn cmd_run(
     ui: &mut Ui,
     command: &CommandHelper,
@@ -708,25 +765,53 @@ pub async fn cmd_run(
     fs::create_dir_all(&base_path)?;
 
     let mut workspace_command = command.workspace_helper(ui).await?;
-    let mut resolved_commits: Vec<_> = if args.revisions.is_empty() {
+    let labeled_revisions: Vec<(RevisionArg, Option<String>)> = if args.revisions.is_empty() {
         let revs = workspace_command.settings().get_string("revsets.run")?;
-        workspace_command
-            .parse_revset(ui, &RevisionArg::from(revs))?
-            .evaluate_to_commits()?
-            .try_collect()
-            .await?
+        vec![parse_labeled_revision(&revs)]
     } else {
-        workspace_command
-            .parse_union_revsets(ui, &args.revisions)?
+        args.revisions
+            .iter()
+            .map(|arg| parse_labeled_revision(arg.as_ref()))
+            .collect()
+    };
+
+    let mut seen_commits: HashMap<CommitId, usize> = HashMap::new();
+    let mut resolved_commits: Vec<(Commit, Option<String>)> = Vec::new();
+    for (rev, label) in &labeled_revisions {
+        let commits: Vec<Commit> = workspace_command
+            .parse_revset(ui, rev)?
             .evaluate_to_commits()?
             .try_collect()
-            .await?
-    };
-    resolved_commits.reverse();
+            .await?;
+        // If a commit is specified in multiple revsets, this is perfectly valid, we
+        // simply take the union of revsets and run it on that. However, if there are
+        // multiple labels, there's no one correct thing to do, so we error out.
+        // If a commit is specified as both labeled and non-labeled, the labeled
+        // wins.
+        for commit in commits.into_iter().rev() {
+            if let Some(&idx) = seen_commits.get(commit.id()) {
+                match (&resolved_commits[idx].1, label) {
+                    (Some(existing), Some(new)) if existing != new => {
+                        return Err(user_error(format!(
+                            "Commit {} has multiple different labels: {existing} and {new}",
+                            workspace_command.format_commit_summary(&commit),
+                        )));
+                    }
+                    (None, Some(new)) => {
+                        resolved_commits[idx].1 = Some(new.clone());
+                    }
+                    _ => {}
+                }
+            } else {
+                seen_commits.insert(commit.id().clone(), resolved_commits.len());
+                resolved_commits.push((commit, label.clone()));
+            }
+        }
+    }
 
     if !args.ignore_changes {
         workspace_command
-            .check_rewritable(resolved_commits.iter().ids())
+            .check_rewritable(resolved_commits.iter().map(|(commit, _)| commit.id()))
             .await?;
     }
 
@@ -861,7 +946,11 @@ pub async fn cmd_run(
     let mut num_reparented: u32 = 0;
     tx.repo_mut()
         .transform_descendants(
-            resolved_commits.iter().ids().cloned().collect_vec(),
+            resolved_commits
+                .iter()
+                .map(|(commit, _)| commit.id())
+                .cloned()
+                .collect_vec(),
             async |rewriter| {
                 let old_id = rewriter.old_commit().id().clone();
                 match (rewritten_commits.get(&old_id), restore_descendants) {
@@ -913,3 +1002,43 @@ pub async fn cmd_run(
 
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_labeled_revision_without_label() {
+        let (rev, label) = parse_labeled_revision("@-");
+        assert_eq!(rev.as_ref(), "@-");
+        assert_eq!(label, None);
+    }
+
+    #[test]
+    fn test_parse_labeled_revision_with_simple_label() {
+        let (rev, label) = parse_labeled_revision("before=@-");
+        assert_eq!(rev.as_ref(), "@-");
+        assert_eq!(label, Some("before".to_string()));
+    }
+
+    #[test]
+    fn test_parse_labeled_revision_with_quoted_string() {
+        let (rev, label) = parse_labeled_revision(r#""branch=name""#);
+        assert_eq!(rev.as_ref(), r#""branch=name""#);
+        assert_eq!(label, None);
+    }
+
+    #[test]
+    fn test_parse_labeled_revision_with_keyword_function() {
+        let (rev, label) = parse_labeled_revision("remote_bookmarks(remote=\"origin\")");
+        assert_eq!(rev.as_ref(), "remote_bookmarks(remote=\"origin\")");
+        assert_eq!(label, None);
+    }
+
+    #[test]
+    fn test_parse_labeled_revision_with_label_and_function() {
+        let (rev, label) = parse_labeled_revision("candidate=remote_bookmarks(remote=\"origin\")");
+        assert_eq!(rev.as_ref(), "remote_bookmarks(remote=\"origin\")");
+        assert_eq!(label, Some("candidate".to_string()));
+    }
+}

cli/tests/cli-reference@.md.snap

@@ -2,6 +2,7 @@
 source: cli/tests/test_generate_md_cli_help.rs
 description: "AUTO-GENERATED FILE, DO NOT EDIT. This cli reference is generated by a test as an `insta` snapshot. MkDocs includes this snapshot from docs/cli-reference.md."
 ---
+
 <!-- BEGIN MARKDOWN-->
 
 # Command-Line Help for `jj`
@@ -3031,6 +3032,7 @@ The command is executed with the following environment variables set:
 - JJ_CHANGE_ID
 - JJ_COMMIT_ID
 - JJ_WORKSPACE_ROOT
+- JJ_LABEL (if a label was specified for the revision)
 
 ### Example
 
@@ -3051,6 +3053,15 @@ $ jj run -j 4 -- pre-commit run .github/pre-commit.yaml
 ###### **Options:**
 
 * `-r`, `--revision <REVSETS>` — The revisions to change
+
+   An optional label prefix can be specified, with at most one label per
+   revision (e.g. `-r before=trunk() -r after=@`). If it is specified:
+
+   * It will be exposed to the command in `$JJ_LABEL`, allowing you to
+     store outputs at known paths - for example: `jj run ... bash -c
+     'run_build && cp out/binary /tmp/$JJ_LABEL'`
+   * It will run in a consistent workspace, ensuring cache affinity for
+     incremental builds.
 * `-j`, `--jobs <JOBS>` — How many processes should run in parallel
 
    Overrides the `run.jobs` config setting. Defaults to 1 if neither is set.