Repository: pantsbuild/pants

Issue:
BSP: Automatically configure the JDK
Currently, after loading a BSP project, a user will likely have to open the project configuration and set the JDK.

<img width="1029" alt="jdk-selection" src="https://user-images.githubusercontent.com/46740/166062813-7f239f43-7362-43aa-a6ed-b161ff067d79.png">

The BSP protocol supports exposing the java home / version of the JDK: https://build-server-protocol.github.io/docs/extensions/jvm.html#jvm-build-target ... and doing so would likely avoid this manual configuration step.


Discussion:

Resolving patch:
diff --git a/src/python/pants/backend/scala/bsp/rules.py b/src/python/pants/backend/scala/bsp/rules.py
--- a/src/python/pants/backend/scala/bsp/rules.py
+++ b/src/python/pants/backend/scala/bsp/rules.py
@@ -3,7 +3,10 @@
 from __future__ import annotations
 
 import logging
+import textwrap
 from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
 
 from pants.backend.scala.bsp.spec import (
     ScalaBuildTarget,
@@ -35,10 +38,20 @@
     BSPResourcesRequest,
     BSPResourcesResult,
 )
+from pants.core.util_rules.system_binaries import BashBinary, ReadlinkBinary, ReadlinkBinaryRequest
 from pants.engine.addresses import Addresses
-from pants.engine.fs import AddPrefix, CreateDigest, Digest, FileEntry, MergeDigests, Workspace
+from pants.engine.fs import (
+    AddPrefix,
+    CreateDigest,
+    Digest,
+    FileContent,
+    FileEntry,
+    MergeDigests,
+    Workspace,
+)
 from pants.engine.internals.native_engine import Snapshot
 from pants.engine.internals.selectors import Get, MultiGet
+from pants.engine.process import Process, ProcessResult
 from pants.engine.rules import _uncacheable_rule, collect_rules, rule
 from pants.engine.target import CoarsenedTarget, CoarsenedTargets, FieldSet, Target, Targets
 from pants.engine.unions import UnionRule
@@ -46,8 +59,9 @@
 from pants.jvm.bsp.compile import rules as jvm_compile_rules
 from pants.jvm.bsp.resources import _jvm_bsp_resources
 from pants.jvm.bsp.resources import rules as jvm_resources_rules
-from pants.jvm.bsp.spec import MavenDependencyModule, MavenDependencyModuleArtifact
+from pants.jvm.bsp.spec import JvmBuildTarget, MavenDependencyModule, MavenDependencyModuleArtifact
 from pants.jvm.compile import ClasspathEntry, ClasspathEntryRequest, ClasspathEntryRequestFactory
+from pants.jvm.jdk_rules import DefaultJdk, JdkEnvironment, JdkRequest
 from pants.jvm.resolve.common import ArtifactRequirement, ArtifactRequirements, Coordinate
 from pants.jvm.resolve.coursier_fetch import (
     CoursierLockfileEntry,
@@ -57,7 +71,8 @@
 )
 from pants.jvm.resolve.key import CoursierResolveKey
 from pants.jvm.subsystems import JvmSubsystem
-from pants.jvm.target_types import JvmArtifactFieldSet, JvmResolveField
+from pants.jvm.target_types import JvmArtifactFieldSet, JvmJdkField, JvmResolveField
+from pants.util.logging import LogLevel
 
 LANGUAGE_ID = "scala"
 
@@ -72,10 +87,11 @@ class ScalaBSPLanguageSupport(BSPLanguageSupport):
 
 @dataclass(frozen=True)
 class ScalaMetadataFieldSet(FieldSet):
-    required_fields = (ScalaSourceField, JvmResolveField)
+    required_fields = (ScalaSourceField, JvmResolveField, JvmJdkField)
 
     source: ScalaSourceField
     resolve: JvmResolveField
+    jdk: JvmJdkField
 
 
 class ScalaBSPResolveFieldFactoryRequest(BSPResolveFieldFactoryRequest):
@@ -154,29 +170,112 @@ def bsp_resolve_field_factory(
 @rule
 async def bsp_resolve_scala_metadata(
     request: ScalaBSPBuildTargetsMetadataRequest,
+    bash: BashBinary,
     jvm: JvmSubsystem,
     scala: ScalaSubsystem,
     build_root: BuildRoot,
 ) -> BSPBuildTargetsMetadataResult:
     resolves = {fs.resolve.normalized_value(jvm) for fs in request.field_sets}
+    jdk_versions = {fs.jdk for fs in request.field_sets}
     if len(resolves) > 1:
         raise ValueError(
             "Cannot provide Scala metadata for multiple resolves. Please set the "
             "`resolve = jvm:$resolve` field in your `[experimental-bsp].groups_config_files` to "
             "select the relevant resolve to use."
         )
-    resolve = list(resolves)[0]
-    scala_version = scala.version_for_resolve(resolve)
+    (resolve,) = resolves
 
+    scala_version = scala.version_for_resolve(resolve)
     scala_runtime = await Get(
         MaterializeScalaRuntimeJarsResult, MaterializeScalaRuntimeJarsRequest(scala_version)
     )
 
+    #
+    # Extract the JDK paths from an lawful-evil process so we can supply it to the IDE.
+    #
+    # Why lawful-evil?
+    # This script relies on implementation details of the Pants JVM execution environment,
+    # namely that the Coursier Archive Cache (i.e. where JDKs are extracted to after download)
+    # is stored into a predictable location on disk and symlinked into the sandbox on process
+    # startup. The script reads the symlink of the cache directory, and outputs the linked
+    # location of the JDK (according to Coursier), and we use that to calculate the permanent
+    # location of the JDK.
+    #
+    # Please don't do anything like this except as a last resort.
+    #
+
+    # The maximum JDK version will be compatible with all the specified targets
+    jdk_requests = [JdkRequest.from_field(version) for version in jdk_versions]
+    jdk_request = max(jdk_requests, key=_jdk_request_sort_key(jvm))
+
+    jdk, readlink, = await MultiGet(
+        Get(JdkEnvironment, JdkRequest, jdk_request),
+        Get(ReadlinkBinary, ReadlinkBinaryRequest()),
+    )
+
+    if any(i.version == DefaultJdk.SYSTEM for i in jdk_requests):
+        system_jdk = await Get(JdkEnvironment, JdkRequest, JdkRequest.SYSTEM)
+        if system_jdk.jre_major_version > jdk.jre_major_version:
+            jdk = system_jdk
+
+    cmd = "leak_paths.sh"
+    leak_jdk_sandbox_paths = textwrap.dedent(
+        f"""\
+        # Script to leak JDK cache paths out of Coursier sandbox so that BSP can use them.
+
+        {readlink.path} {jdk.coursier.cache_dir}
+        {jdk.java_home_command}
+        """
+    )
+    leak_sandbox_path_digest = await Get(
+        Digest,
+        CreateDigest(
+            [
+                FileContent(
+                    cmd,
+                    leak_jdk_sandbox_paths.encode("utf-8"),
+                    is_executable=True,
+                ),
+            ]
+        ),
+    )
+
+    leaked_paths = await Get(
+        ProcessResult,
+        Process(
+            [
+                bash.path,
+                cmd,
+            ],
+            input_digest=leak_sandbox_path_digest,
+            immutable_input_digests=jdk.immutable_input_digests,
+            env=jdk.env,
+            use_nailgun=(),
+            description="Report JDK cache paths for BSP",
+            append_only_caches=jdk.append_only_caches,
+            level=LogLevel.TRACE,
+        ),
+    )
+
+    cache_dir, jdk_home = leaked_paths.stdout.decode().strip().split("\n")
+
+    _, sep, suffix = jdk_home.partition(jdk.coursier.cache_dir)
+    if sep:
+        coursier_java_home = cache_dir + suffix
+    else:
+        # Partition failed. Probably a system JDK instead
+        coursier_java_home = jdk_home
+
     scala_jar_uris = tuple(
         build_root.pathlib_path.joinpath(".pants.d/bsp").joinpath(p).as_uri()
         for p in scala_runtime.content.files
     )
 
+    jvm_build_target = JvmBuildTarget(
+        java_home=Path(coursier_java_home).as_uri(),
+        java_version=f"1.{jdk.jre_major_version}",
+    )
+
     return BSPBuildTargetsMetadataResult(
         metadata=ScalaBuildTarget(
             scala_organization="org.scala-lang",
@@ -184,11 +283,27 @@ async def bsp_resolve_scala_metadata(
             scala_binary_version=".".join(scala_version.split(".")[0:2]),
             platform=ScalaPlatform.JVM,
             jars=scala_jar_uris,
+            jvm_build_target=jvm_build_target,
         ),
         digest=scala_runtime.content.digest,
     )
 
 
+def _jdk_request_sort_key(
+    jvm: JvmSubsystem,
+) -> Callable[[JdkRequest,], tuple[int, ...]]:
+    def sort_key_function(request: JdkRequest) -> tuple[int, ...]:
+        if request == JdkRequest.SYSTEM:
+            return (-1,)
+
+        version_str = request.version if isinstance(request.version, str) else jvm.jdk
+        _, version = version_str.split(":")
+
+        return tuple(int(i) for i in version.split("."))
+
+    return sort_key_function
+
+
 # -----------------------------------------------------------------------------------------------
 # Scalac Options Request
 # See https://build-server-protocol.github.io/docs/extensions/scala.html#scalac-options-request
diff --git a/src/python/pants/core/util_rules/system_binaries.py b/src/python/pants/core/util_rules/system_binaries.py
--- a/src/python/pants/core/util_rules/system_binaries.py
+++ b/src/python/pants/core/util_rules/system_binaries.py
@@ -327,6 +327,10 @@ class DiffBinary(BinaryPath):
     pass
 
 
+class ReadlinkBinary(BinaryPath):
+    pass
+
+
 class GitBinaryException(Exception):
     pass
 
@@ -690,6 +694,14 @@ async def find_diff() -> DiffBinary:
     return DiffBinary(first_path.path, first_path.fingerprint)
 
 
+@rule(desc="Finding the `readlink` binary", level=LogLevel.DEBUG)
+async def find_readlink() -> ReadlinkBinary:
+    request = BinaryPathRequest(binary_name="readlink", search_path=SEARCH_PATHS)
+    paths = await Get(BinaryPaths, BinaryPathRequest, request)
+    first_path = paths.first_path_or_raise(request, rationale="defererence symlinks")
+    return ReadlinkBinary(first_path.path, first_path.fingerprint)
+
+
 @rule(desc="Finding the `git` binary", level=LogLevel.DEBUG)
 async def find_git() -> GitBinary:
     request = BinaryPathRequest(binary_name="git", search_path=SEARCH_PATHS)
@@ -734,6 +746,10 @@ class DiffBinaryRequest:
     pass
 
 
+class ReadlinkBinaryRequest:
+    pass
+
+
 class GitBinaryRequest:
     pass
 
@@ -763,6 +779,13 @@ async def find_mkdir_wrapper(_: MkdirBinaryRequest, mkdir_binary: MkdirBinary) -
     return mkdir_binary
 
 
+@rule
+async def find_readlink_wrapper(
+    _: ReadlinkBinaryRequest, readlink_binary: ReadlinkBinary
+) -> ReadlinkBinary:
+    return readlink_binary
+
+
 @rule
 async def find_chmod_wrapper(_: ChmodBinaryRequest, chmod_binary: ChmodBinary) -> ChmodBinary:
     return chmod_binary
diff --git a/src/python/pants/jvm/jdk_rules.py b/src/python/pants/jvm/jdk_rules.py
--- a/src/python/pants/jvm/jdk_rules.py
+++ b/src/python/pants/jvm/jdk_rules.py
@@ -95,6 +95,7 @@ class JdkEnvironment:
     coursier: Coursier
     jre_major_version: int
     global_jvm_options: tuple[str, ...]
+    java_home_command: str
 
     bin_dir: ClassVar[str] = "__jdk"
     jdk_preparation_script: ClassVar[str] = f"{bin_dir}/jdk.sh"
@@ -138,6 +139,7 @@ def from_jdk_environment(cls, env: JdkEnvironment) -> InternalJdk:
             env.coursier,
             env.jre_major_version,
             env.global_jvm_options,
+            env.java_home_command,
         )
 
 
@@ -293,6 +295,7 @@ def prefixed(arg: str) -> str:
         nailgun_jar=os.path.join(JdkEnvironment.bin_dir, nailgun.filenames[0]),
         coursier=coursier,
         jre_major_version=jre_major_version,
+        java_home_command=java_home_command,
     )
 
 


Test patch: