[internal] jvm/java: ensure JDK downloaded in one process (#12972)

## Motivation

As described in #12293, multiple Coursier invocations were downloading the JDK and triggering [a race condition in Coursier's locking](coursier/coursier#1815) that caused flakiness in tests.

## Solution

This PR mitigates the issue by isolating JDK download to a single `Process`. The new `JdkSetup` type provides rules with the command to obtain the location of the JDK so they may query Coursier for JAVA_HOME. This has the benefit of still downloading in remote execution, but providing some guarantee that there will be a single download.

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

@@ -2,11 +2,12 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 from pants.backend.java import tailor
+from pants.backend.java import util_rules as java_util_rules
 from pants.backend.java.compile import javac, javac_binary
 from pants.backend.java.target_types import JavaLibrary, JunitTests
 from pants.backend.java.target_types import rules as target_types_rules
 from pants.backend.java.test import junit
-from pants.jvm import util_rules
+from pants.jvm import util_rules as jvm_util_rules
 from pants.jvm.goals import coursier
 from pants.jvm.resolve import coursier_fetch, coursier_setup
 from pants.jvm.target_types import JvmDependencyLockfile
@@ -25,6 +26,7 @@ def rules():
         *coursier_fetch.rules(),
         *coursier_setup.rules(),
         *tailor.rules(),
-        *util_rules.rules(),
+        *jvm_util_rules.rules(),
+        *java_util_rules.rules(),
         *target_types_rules(),
     ]

src/python/pants/backend/java/compile/javac_binary.py

@@ -8,11 +8,9 @@
 from dataclasses import dataclass
 from typing import ClassVar
 
-from pants.backend.java.compile.javac_subsystem import JavacSubsystem
+from pants.backend.java.util_rules import JdkSetup
 from pants.engine.fs import CreateDigest, Digest, FileContent, MergeDigests
-from pants.engine.process import Process, ProcessCacheScope, ProcessResult
 from pants.engine.rules import Get, collect_rules, rule
-from pants.jvm.resolve.coursier_setup import Coursier
 
 logger = logging.getLogger(__name__)
 
@@ -25,52 +23,16 @@ class JavacBinary:
 
 
 @rule
-async def setup_javac_binary(coursier: Coursier, javac: JavacSubsystem) -> JavacBinary:
-    if javac.options.jdk == "system":
-        process_result = await Get(
-            ProcessResult,
-            Process(
-                argv=[
-                    coursier.coursier.exe,
-                    "java",
-                    "--system-jvm",
-                    "-version",
-                ],
-                input_digest=coursier.digest,
-                description="Invoke Coursier with system-jvm to fingerprint JVM version.",
-                cache_scope=ProcessCacheScope.PER_RESTART_SUCCESSFUL,
-            ),
-        )
-        all_output = "\n".join(
-            [
-                process_result.stderr.decode("utf-8"),
-                process_result.stdout.decode("utf-8"),
-            ]
-        )
-        fingerprint_comment_lines = [
-            "pants javac script using Coursier --system-jvm.  System java -version:",
-            *filter(None, all_output.splitlines()),
-        ]
-        fingerprint_comment = "".join([f"# {line}\n" for line in fingerprint_comment_lines])
-        javac_path_line = (
-            f'javac_path="$({coursier.coursier.exe} java-home --system-jvm)/bin/javac"'
-        )
-    else:
-        fingerprint_comment = f"# pants javac script using Coursier with --jvm {javac.options.jdk}"
-        javac_path_line = (
-            f'javac_path="$({coursier.coursier.exe} java-home --jvm {javac.options.jdk})/bin/javac"'
-        )
-
+async def setup_javac_binary(jdk_setup: JdkSetup) -> JavacBinary:
     # Awkward join so multi-line `fingerprint_comment` won't confuse textwrap.dedent
     javac_wrapper_script = "\n".join(
         [
-            fingerprint_comment,
+            jdk_setup.fingerprint_comment,
             textwrap.dedent(
                 f"""\
                 set -eu
-                {javac_path_line}
                 /bin/mkdir -p {JavacBinary.classfiles_relpath}
-                exec "${{javac_path}}" "$@"
+                exec $({' '.join(jdk_setup.java_home_cmd)})/bin/javac "$@"
                 """
             ),
         ]
@@ -92,7 +54,7 @@ async def setup_javac_binary(coursier: Coursier, javac: JavacSubsystem) -> Javac
             Digest,
             MergeDigests(
                 [
-                    coursier.digest,
+                    jdk_setup.digest,
                     javac_wrapper_script_digest,
                 ]
             ),

src/python/pants/backend/java/compile/javac_binary_test.py

@@ -7,6 +7,7 @@
 
 from pants.backend.java.compile.javac_binary import JavacBinary
 from pants.backend.java.compile.javac_binary import rules as javac_binary_rules
+from pants.backend.java.util_rules import rules as util_rules_rules
 from pants.core.util_rules.external_tool import rules as external_tool_rules
 from pants.engine.internals.scheduler import ExecutionError
 from pants.engine.process import BashBinary, Process, ProcessResult
@@ -22,6 +23,7 @@ def rule_runner() -> RuleRunner:
             *coursier_setup_rules(),
             *external_tool_rules(),
             *javac_binary_rules(),
+            *util_rules_rules(),
             *process_rules(),
             QueryRule(BashBinary, ()),
             QueryRule(JavacBinary, ()),

src/python/pants/backend/java/compile/javac_test.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from pants.backend.java import util_rules as java_util_rules
 from pants.backend.java.compile.javac import (
     CompiledClassfiles,
     CompileJavaSourceRequest,
@@ -54,6 +55,7 @@ def rule_runner() -> RuleRunner:
             *javac_binary_rules(),
             *target_types_rules(),
             *coursier_rules(),
+            *java_util_rules.rules(),
             QueryRule(CheckResults, (JavacCheckRequest,)),
             QueryRule(FallibleCompiledClassfiles, (CompileJavaSourceRequest,)),
             QueryRule(CompiledClassfiles, (CompileJavaSourceRequest,)),

src/python/pants/backend/java/dependency_inference/java_parser.py

@@ -13,12 +13,12 @@
     java_parser_artifact_requirements,
 )
 from pants.backend.java.dependency_inference.types import JavaSourceDependencyAnalysis
+from pants.backend.java.util_rules import JdkSetup
 from pants.core.util_rules.source_files import SourceFiles
 from pants.engine.fs import AddPrefix, Digest, DigestContents, MergeDigests
 from pants.engine.process import BashBinary, FallibleProcessResult, Process, ProcessExecutionFailure
 from pants.engine.rules import Get, collect_rules, rule
 from pants.jvm.resolve.coursier_fetch import MaterializedClasspath, MaterializedClasspathRequest
-from pants.jvm.resolve.coursier_setup import Coursier
 from pants.util.logging import LogLevel
 
 logger = logging.getLogger(__name__)
@@ -52,7 +52,7 @@ async def resolve_fallible_result_to_analysis(
 @rule(level=LogLevel.DEBUG)
 async def analyze_java_source_dependencies(
     bash: BashBinary,
-    coursier: Coursier,
+    jdk_setup: JdkSetup,
     processor_classfiles: JavaParserCompiledClassfiles,
     source_files: SourceFiles,
 ) -> FallibleJavaSourceDependencyAnalysisResult:
@@ -88,8 +88,8 @@ async def analyze_java_source_dependencies(
             (
                 prefixed_processor_classfiles_digest,
                 tool_classpath.digest,
-                coursier.digest,
                 prefixed_source_files_digest,
+                jdk_setup.digest,
             )
         ),
     )
@@ -98,9 +98,10 @@ async def analyze_java_source_dependencies(
 
     proc = Process(
         argv=[
-            coursier.coursier.exe,
-            "java",
-            "--system-jvm",  # TODO(#12293): use a fixed JDK version from a subsystem.
+            bash.path,
+            "-c",
+            f"exec $({' '.join(jdk_setup.java_home_cmd)})/bin/java \"$@\"",
+            "--",
             "-cp",
             ":".join([tool_classpath.classpath_arg(), processorcp_relpath]),
             "org.pantsbuild.javaparser.PantsJavaParserLauncher",

src/python/pants/backend/java/dependency_inference/java_parser_test.py

@@ -7,6 +7,7 @@
 
 import pytest
 
+from pants.backend.java import util_rules as java_util_rules
 from pants.backend.java.compile.javac import rules as javac_rules
 from pants.backend.java.compile.javac_binary import rules as javac_binary_rules
 from pants.backend.java.dependency_inference.java_parser import (
@@ -47,6 +48,7 @@ def rule_runner() -> RuleRunner:
             *javac_rules(),
             *source_files.rules(),
             *util_rules(),
+            *java_util_rules.rules(),
             QueryRule(FallibleJavaSourceDependencyAnalysisResult, (SourceFiles,)),
             QueryRule(JavaSourceDependencyAnalysis, (SourceFiles,)),
             QueryRule(SourceFiles, (SourceFilesRequest,)),

src/python/pants/backend/java/test/junit.py

@@ -7,10 +7,11 @@
 from pants.backend.java.compile.javac import CompiledClassfiles, CompileJavaSourceRequest
 from pants.backend.java.subsystems.junit import JUnit
 from pants.backend.java.target_types import JavaTestsSources
+from pants.backend.java.util_rules import JdkSetup
 from pants.core.goals.test import TestDebugRequest, TestFieldSet, TestResult, TestSubsystem
 from pants.engine.addresses import Addresses
 from pants.engine.fs import AddPrefix, Digest, MergeDigests
-from pants.engine.process import FallibleProcessResult, Process
+from pants.engine.process import BashBinary, FallibleProcessResult, Process
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import (
     CoarsenedTargets,
@@ -27,7 +28,6 @@
     MaterializedClasspath,
     MaterializedClasspathRequest,
 )
-from pants.jvm.resolve.coursier_setup import Coursier
 from pants.util.logging import LogLevel
 
 logger = logging.getLogger(__name__)
@@ -42,7 +42,8 @@ class JavaTestFieldSet(TestFieldSet):
 
 @rule(desc="Run JUnit", level=LogLevel.DEBUG)
 async def run_junit_test(
-    coursier: Coursier,
+    bash: BashBinary,
+    jdk_setup: JdkSetup,
     junit: JUnit,
     test_subsystem: TestSubsystem,
     field_set: JavaTestFieldSet,
@@ -99,15 +100,16 @@ async def run_junit_test(
             (
                 prefixed_transitive_user_classfiles_digest,
                 materialized_classpath.digest,
-                coursier.digest,
+                jdk_setup.digest,
             )
         ),
     )
     proc = Process(
         argv=[
-            coursier.coursier.exe,
-            "java",
-            "--system-jvm",  # TODO(#12293): use a fixed JDK version from a subsystem.
+            bash.path,
+            "-c",
+            f"exec $({' '.join(jdk_setup.java_home_cmd)})/bin/java \"$@\"",
+            "--",
             "-cp",
             materialized_classpath.classpath_arg(),
             "org.junit.platform.console.ConsoleLauncher",

src/python/pants/backend/java/test/junit_test.py

@@ -14,6 +14,7 @@
 from pants.backend.java.target_types import rules as target_types_rules
 from pants.backend.java.test.junit import JavaTestFieldSet
 from pants.backend.java.test.junit import rules as junit_rules
+from pants.backend.java.util_rules import rules as java_util_rules
 from pants.build_graph.address import Address
 from pants.core.goals.test import TestResult
 from pants.core.util_rules import config_files, source_files
@@ -50,6 +51,7 @@ def rule_runner() -> RuleRunner:
             *junit_rules(),
             *javac_binary_rules(),
             *util_rules(),
+            *java_util_rules(),
             *target_types_rules(),
             QueryRule(CoarsenedTargets, (Addresses,)),
             QueryRule(TestResult, (JavaTestFieldSet,)),

src/python/pants/backend/java/util_rules.py

@@ -0,0 +1,81 @@
+# Copyright 2021 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pants.backend.java.compile.javac_subsystem import JavacSubsystem
+from pants.engine.fs import Digest
+from pants.engine.internals.selectors import Get
+from pants.engine.process import FallibleProcessResult, Process, ProcessResult
+from pants.engine.rules import collect_rules, rule
+from pants.jvm.resolve.coursier_setup import Coursier
+
+
+@dataclass(frozen=True)
+class JdkSetup:
+    digest: Digest
+    java_home_cmd: tuple[str, ...]
+    fingerprint_comment: str
+
+
+@rule
+async def setup_jdk(coursier: Coursier, javac: JavacSubsystem) -> JdkSetup:
+    if javac.options.jdk == "system":
+        coursier_jdk_option = "--system-jvm"
+    else:
+        coursier_jdk_option = f"--jvm={javac.options.jdk}"
+
+    java_home_result = await Get(
+        FallibleProcessResult,
+        Process(
+            argv=(
+                coursier.coursier.exe,
+                "java-home",
+                coursier_jdk_option,
+            ),
+            input_digest=coursier.digest,
+            description=f"Ensure download of JDK {coursier_jdk_option}.",
+        ),
+    )
+
+    if java_home_result.exit_code != 0:
+        raise ValueError(
+            f"Failed to determine JAVA_HOME for JDK {javac.options.jdk}: {java_home_result.stderr.decode('utf-8')}"
+        )
+
+    java_home = java_home_result.stdout.decode("utf-8").strip()
+
+    version_result = await Get(
+        ProcessResult,
+        Process(
+            argv=(
+                f"{java_home}/bin/java",
+                "-version",
+            ),
+            description=f"Extract version from JDK {coursier_jdk_option}.",
+        ),
+    )
+
+    all_output = "\n".join(
+        [
+            version_result.stderr.decode("utf-8"),
+            version_result.stdout.decode("utf-8"),
+        ]
+    )
+    fingerprint_comment_lines = [
+        f"pants javac script using Coursier {coursier_jdk_option}.  `java -version`:",
+        *filter(None, all_output.splitlines()),
+    ]
+    fingerprint_comment = "".join([f"# {line}\n" for line in fingerprint_comment_lines])
+
+    return JdkSetup(
+        digest=coursier.digest,
+        java_home_cmd=(coursier.coursier.exe, "java-home", coursier_jdk_option),
+        fingerprint_comment=fingerprint_comment,
+    )
+
+
+def rules():
+    return collect_rules()