Adds run support for deploy_jar targets (pantsbuild#14352)

This adds `run` support for `deploy_jar` targets. It works! Closes pantsbuild#14283, at least until `war` files are properly supported.

To make Scala binaries work, we now inject a dependency for the version of `scala-library` specified by `--scala-version` at compile time, whenever `scala-library` is not provided as a transitive dependency. This is a good-enough stop-gap in service of pantsbuild#14171.

There's one ugly caveat which I'm reserving for future work:

For tool support, Couriser downloads a JVM into an append-only cache for most JVM processes. Since `InteractiveProcess` does not support these, I've added `RuntimeJdk__`, which downloads a JDK into a `Digest`. The digest gets written into the temporary directory alongside the runnable artifacts. It's not ideal, but for now, it removes the JDK download process from the `run` process, and the JDK is only downloaded once per `pantsd` execution. It's probably _ok_ for now. Being able to reuse the global append-only cache directories will be a better solution once we get there.

src/python/pants/backend/experimental/java/register.py

@@ -12,7 +12,7 @@
     JunitTestTarget,
 )
 from pants.backend.java.target_types import rules as target_types_rules
-from pants.jvm import classpath, jdk_rules, resources
+from pants.jvm import classpath, jdk_rules, resources, run_deploy_jar
 from pants.jvm import util_rules as jvm_util_rules
 from pants.jvm.dependency_inference import symbol_mapper
 from pants.jvm.goals import lockfile
@@ -52,4 +52,5 @@ def rules():
         *jdk_rules.rules(),
         *target_types_rules(),
         *jvm_tool.rules(),
+        *run_deploy_jar.rules(),
     ]

src/python/pants/backend/experimental/scala/register.py

@@ -14,7 +14,7 @@
 )
 from pants.backend.scala.target_types import rules as target_types_rules
 from pants.backend.scala.test import scalatest
-from pants.jvm import classpath, jdk_rules, resources
+from pants.jvm import classpath, jdk_rules, resources, run_deploy_jar
 from pants.jvm import util_rules as jvm_util_rules
 from pants.jvm.goals import lockfile
 from pants.jvm.package import deploy_jar
@@ -56,4 +56,5 @@ def rules():
         *target_types_rules(),
         *jvm_tool.rules(),
         *resources.rules(),
+        *run_deploy_jar.rules(),
     ]

src/python/pants/backend/scala/compile/scalac.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 from itertools import chain
 
 from pants.backend.java.target_types import JavaFieldSet, JavaGeneratorFieldSet, JavaSourceField
@@ -39,6 +40,11 @@ class CompileScalaSourceRequest(ClasspathEntryRequest):
     field_sets_consume_only = (JavaFieldSet, JavaGeneratorFieldSet)
 
 
+@dataclass(frozen=True)
+class ScalaLibraryRequest:
+    version: str
+
+
 @rule(desc="Compile with scalac")
 async def compile_scala_source(
     scala: ScalaSubsystem,
@@ -68,6 +74,20 @@ async def compile_scala_source(
             exit_code=1,
         )
 
+    all_dependency_jars = [
+        filename
+        for dependency in direct_dependency_classpath_entries
+        for filename in dependency.filenames
+    ]
+
+    # TODO(14171): Stop-gap for making sure that scala supplies all of its dependencies to
+    # deploy targets.
+    if not any(
+        filename.startswith("org.scala-lang_scala-library_") for filename in all_dependency_jars
+    ):
+        scala_library = await Get(ClasspathEntry, ScalaLibraryRequest(scala.version))
+        direct_dependency_classpath_entries += (scala_library,)
+
     component_members_with_sources = tuple(
         t for t in request.component.members if t.has_field(SourcesField)
     )
@@ -189,6 +209,26 @@ async def compile_scala_source(
     )
 
 
+@rule
+async def fetch_scala_library(request: ScalaLibraryRequest) -> ClasspathEntry:
+    tcp = await Get(
+        ToolClasspath,
+        ToolClasspathRequest(
+            artifact_requirements=ArtifactRequirements.from_coordinates(
+                [
+                    Coordinate(
+                        group="org.scala-lang",
+                        artifact="scala-library",
+                        version=request.version,
+                    ),
+                ]
+            ),
+        ),
+    )
+
+    return ClasspathEntry(tcp.digest, tcp.content.files)
+
+
 def rules():
     return [
         *collect_rules(),

src/python/pants/jvm/compile_test.py

@@ -216,16 +216,19 @@ def test_compile_mixed(rule_runner: RuleRunner) -> None:
     rendered_classpath = rule_runner.request(
         RenderedClasspath, [Addresses([Address(spec_path="", target_name="main")])]
     )
-    assert rendered_classpath.content == {
-        ".Example.scala.main.scalac.jar": {
-            "META-INF/MANIFEST.MF",
-            "org/pantsbuild/example/Main$.class",
-            "org/pantsbuild/example/Main.class",
-        },
-        "lib.C.java.javac.jar": {
-            "org/pantsbuild/example/lib/C.class",
-        },
+
+    assert rendered_classpath.content[".Example.scala.main.scalac.jar"] == {
+        "META-INF/MANIFEST.MF",
+        "org/pantsbuild/example/Main$.class",
+        "org/pantsbuild/example/Main.class",
+    }
+    assert rendered_classpath.content["lib.C.java.javac.jar"] == {
+        "org/pantsbuild/example/lib/C.class",
     }
+    assert any(
+        key.startswith("org.scala-lang_scala-library_") for key in rendered_classpath.content.keys()
+    )
+    assert len(rendered_classpath.content.keys()) == 3
 
 
 @maybe_skip_jdk_test

src/python/pants/jvm/jdk_rules.py

@@ -91,9 +91,26 @@ async def setup_jdk(coursier: Coursier, jvm: JvmSubsystem, bash: BashBinary) ->
         coursier_jdk_option = "--system-jvm"
     else:
         coursier_jdk_option = shlex.quote(f"--jvm={jvm.jdk}")
+
+    # TODO(#14386) This argument re-writing code should be done in a more standardised way.
+    # See also `run_deploy_jar` for other argument re-writing code.
+    def prefixed(arg: str) -> str:
+        if arg.startswith("__"):
+            return f"${{PANTS_INTERNAL_ABSOLUTE_PREFIX}}{arg}"
+        else:
+            return arg
+
+    optionally_prefixed_coursier_args = [
+        prefixed(arg) for arg in coursier.args(["java-home", coursier_jdk_option])
+    ]
     # NB: We `set +e` in the subshell to ensure that it exits as well.
     #  see https://unix.stackexchange.com/a/23099
-    java_home_command = " ".join(("set +e;", *coursier.args(["java-home", coursier_jdk_option])))
+    java_home_command = " ".join(("set +e;", *optionally_prefixed_coursier_args))
+
+    env = {
+        "PANTS_INTERNAL_ABSOLUTE_PREFIX": "",
+        **coursier.env,
+    }
 
     java_version_result = await Get(
         FallibleProcessResult,
@@ -105,7 +122,7 @@ async def setup_jdk(coursier: Coursier, jvm: JvmSubsystem, bash: BashBinary) ->
             ),
             append_only_caches=coursier.append_only_caches,
             immutable_input_digests=coursier.immutable_input_digests,
-            env=coursier.env,
+            env=env,
             description=f"Ensure download of JDK {coursier_jdk_option}.",
             cache_scope=ProcessCacheScope.PER_RESTART_SUCCESSFUL,
             level=LogLevel.DEBUG,
@@ -135,7 +152,7 @@ async def setup_jdk(coursier: Coursier, jvm: JvmSubsystem, bash: BashBinary) ->
         {version_comment}
         set -eu
 
-        /bin/ln -s "$({java_home_command})" "{JdkSetup.java_home}"
+        /bin/ln -s "$({java_home_command})" "${{PANTS_INTERNAL_ABSOLUTE_PREFIX}}{JdkSetup.java_home}"
         exec "$@"
         """
     )
@@ -232,7 +249,11 @@ async def jvm_process(bash: BashBinary, jdk_setup: JdkSetup, request: JvmProcess
         **jdk_setup.immutable_input_digests,
         **request.extra_immutable_input_digests,
     }
-    env = {**jdk_setup.env, **request.extra_env}
+    env = {
+        "PANTS_INTERNAL_ABSOLUTE_PREFIX": "",
+        **jdk_setup.env,
+        **request.extra_env,
+    }
 
     use_nailgun = []
     if request.use_nailgun:

src/python/pants/jvm/package/deploy_jar.py

@@ -13,6 +13,7 @@
     OutputPathField,
     PackageFieldSet,
 )
+from pants.core.goals.run import RunFieldSet
 from pants.core.util_rules.archive import ZipBinary
 from pants.engine.addresses import Addresses
 from pants.engine.fs import EMPTY_DIGEST, AddPrefix, CreateDigest, Digest, FileContent, MergeDigests
@@ -39,7 +40,7 @@
 
 
 @dataclass(frozen=True)
-class DeployJarFieldSet(PackageFieldSet):
+class DeployJarFieldSet(PackageFieldSet, RunFieldSet):
     required_fields = (
         JvmMainClassNameField,
         Dependencies,

src/python/pants/jvm/run_deploy_jar.py

@@ -0,0 +1,138 @@
+# Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import dataclasses
+import logging
+from dataclasses import dataclass
+from typing import Iterable
+
+from pants.core.goals.package import BuiltPackage
+from pants.core.goals.run import RunFieldSet, RunRequest
+from pants.engine.fs import EMPTY_DIGEST, Digest, MergeDigests
+from pants.engine.internals.native_engine import AddPrefix
+from pants.engine.process import Process, ProcessResult
+from pants.engine.rules import Get, MultiGet, collect_rules, rule
+from pants.engine.unions import UnionRule
+from pants.jvm.jdk_rules import JdkSetup, JvmProcess
+from pants.jvm.package.deploy_jar import DeployJarFieldSet
+from pants.util.frozendict import FrozenDict
+from pants.util.logging import LogLevel
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class __RuntimeJvm:
+    """Allows Coursier to download a JDK into a Digest, rather than an append-only cache for use
+    with `pants run`.
+
+    This is a hideous stop-gap, which will no longer be necessary once `InteractiveProcess` supports
+    append-only caches. (See #13852 for details on how to do this.)
+    """
+
+    digest: Digest
+
+
+@rule(level=LogLevel.DEBUG)
+async def create_deploy_jar_run_request(
+    field_set: DeployJarFieldSet,
+    runtime_jvm: __RuntimeJvm,
+    jdk_setup: JdkSetup,
+) -> RunRequest:
+
+    main_class = field_set.main_class.value
+    assert main_class is not None
+
+    package = await Get(BuiltPackage, DeployJarFieldSet, field_set)
+    assert len(package.artifacts) == 1
+    jar_path = package.artifacts[0].relpath
+    assert jar_path is not None
+
+    proc = await Get(
+        Process,
+        JvmProcess(
+            classpath_entries=[f"{{chroot}}/{jar_path}"],
+            argv=(main_class,),
+            input_digest=package.digest,
+            description=f"Run {main_class}.main(String[])",
+            use_nailgun=False,
+        ),
+    )
+
+    support_digests = await MultiGet(
+        Get(Digest, AddPrefix(digest, prefix))
+        for prefix, digest in proc.immutable_input_digests.items()
+    )
+
+    support_digests += (runtime_jvm.digest,)
+
+    # TODO(#14386) This argument re-writing code should be done in a more standardised way.
+    # See also `jdk_rules.py` for other argument re-writing code.
+    def prefixed(arg: str, prefixes: Iterable[str]) -> str:
+        if any(arg.startswith(prefix) for prefix in prefixes):
+            return f"{{chroot}}/{arg}"
+        else:
+            return arg
+
+    prefixes = (jdk_setup.bin_dir, jdk_setup.jdk_preparation_script, jdk_setup.java_home)
+    args = [prefixed(arg, prefixes) for arg in proc.argv]
+
+    env = {
+        **proc.env,
+        "PANTS_INTERNAL_ABSOLUTE_PREFIX": "{chroot}/",
+    }
+
+    # absolutify coursier cache envvars
+    for key in env:
+        if key.startswith("COURSIER"):
+            env[key] = prefixed(env[key], (jdk_setup.coursier.cache_dir,))
+
+    request_digest = await Get(
+        Digest,
+        MergeDigests(
+            [
+                proc.input_digest,
+                *support_digests,
+            ]
+        ),
+    )
+
+    return RunRequest(
+        digest=request_digest,
+        args=args,
+        extra_env=env,
+    )
+
+
+@rule
+async def ensure_jdk_for_pants_run(jdk_setup: JdkSetup) -> __RuntimeJvm:
+    # `tools.jar` is distributed with the JDK, so we can rely on it existing.
+    ensure_jvm_process = await Get(
+        Process,
+        JvmProcess(
+            classpath_entries=[f"{jdk_setup.java_home}/lib/tools.jar"],
+            argv=["com.sun.tools.javac.Main", "--version"],
+            input_digest=EMPTY_DIGEST,
+            description="Ensure download of JDK for `pants run` use",
+        ),
+    )
+
+    # Do not treat the coursier JDK digest an append-only cache, so that we can capture the
+    # downloaded JDK in a `Digest`
+    new_append_only_caches = {
+        "coursier_archive": ".cache/arc",
+        "coursier_jvm": ".cache/jvm",
+    }
+
+    ensure_jvm_process = dataclasses.replace(
+        ensure_jvm_process,
+        append_only_caches=FrozenDict(new_append_only_caches),
+        output_directories=(".cache/jdk",),
+    )
+    ensure_jvm = await Get(ProcessResult, Process, ensure_jvm_process)
+
+    return __RuntimeJvm(ensure_jvm.output_digest)
+
+
+def rules():
+    return [*collect_rules(), UnionRule(RunFieldSet, DeployJarFieldSet)]