[fix] Compute classpath lazily in JvmEnvironment requests

Merge-request: BAZEL-MR-658 Merged-by: Tomasz Pasternak <Tomasz.Pasternak@jetbrains.com>
JetBrains · Nov 23, 2023 · f6823e7 · f6823e7
1 parent 24aa30d
commit f6823e7
Show file tree

Hide file tree

Showing 11 changed files with 74 additions and 33 deletions.
diff --git a/aspects/runtime_classpath_query.bzl b/aspects/runtime_classpath_query.bzl
@@ -0,0 +1,10 @@
+def format(target):
+    provider = providers(target)["JavaInfo"]
+    compilation_info = getattr(provider, "compilation_info", None)
+
+    result = []
+    if compilation_info:
+        result = compilation_info.runtime_classpath.to_list()
+    elif hasattr(provider, "transitive_runtime_jars"):
+        result = provider.transitive_runtime_jars.to_list()
+    return "\n".join([f.path for f in result])
diff --git a/bazelrunner/src/main/kotlin/org/jetbrains/bsp/bazel/bazelrunner/BazelRunnerCommandBuilder.kt b/bazelrunner/src/main/kotlin/org/jetbrains/bsp/bazel/bazelrunner/BazelRunnerCommandBuilder.kt
@@ -15,6 +15,8 @@ class BazelRunnerCommandBuilder internal constructor(private val bazelRunner: Ba
   fun showRepo() = mod("show_repo")
   fun showExtension() = mod("show_extension")
   fun query()  = BazelRunnerBuilder(bazelRunner, listOf("query"))
+  fun cquery()  = BazelRunnerBuilder(bazelRunner, listOf("cquery"))
+
   fun build()  = BazelRunnerBuildBuilder(bazelRunner, listOf("build")).withUseBuildFlags()
   fun test()   = BazelRunnerBuildBuilder(bazelRunner, listOf("test")).withUseBuildFlags()
 }
diff --git a/bspcli/src/main/kotlin/org/jetbrains/bsp/cli/Main.kt b/bspcli/src/main/kotlin/org/jetbrains/bsp/cli/Main.kt
@@ -7,6 +7,7 @@ import ch.epfl.scala.bsp4j.DidChangeBuildTarget
 import ch.epfl.scala.bsp4j.InitializeBuildParams
 import ch.epfl.scala.bsp4j.JavaBuildServer
 import ch.epfl.scala.bsp4j.JvmBuildServer
+import ch.epfl.scala.bsp4j.JvmRunEnvironmentParams
 import ch.epfl.scala.bsp4j.LogMessageParams
 import ch.epfl.scala.bsp4j.PrintParams
 import ch.epfl.scala.bsp4j.PublishDiagnosticsParams

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/BazelBspServer.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/BazelBspServer.kt
@@ -47,7 +47,11 @@ class BazelBspServer(
         )
 
         val bspProjectMapper = BspProjectMapper(
-            serverContainer.languagePluginsService, workspaceContextProvider
+                serverContainer.languagePluginsService,
+                workspaceContextProvider,
+                serverContainer.bazelPathsResolver,
+                bazelRunner,
+                bspInfo
         )
         val projectSyncService =
             ProjectSyncService(bspProjectMapper, serverContainer.projectProvider)

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/common/ServerContainer.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/common/ServerContainer.kt
@@ -40,7 +40,8 @@ class ServerContainer internal constructor(
   val bazelInfo: BazelInfo,
   val bazelRunner: BazelRunner,
   val compilationManager: BazelBspCompilationManager,
-  val languagePluginsService: LanguagePluginsService
+  val languagePluginsService: LanguagePluginsService,
+  val bazelPathsResolver: BazelPathsResolver,
 ) {
   companion object {
     @JvmStatic
@@ -105,7 +106,8 @@ class ServerContainer internal constructor(
         bazelInfo,
         bazelRunner,
         compilationManager,
-        languagePluginsService
+        languagePluginsService,
+        bazelPathsResolver,
       )
     }
   }

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BazelPathsResolver.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BazelPathsResolver.kt
@@ -17,6 +17,8 @@ class BazelPathsResolver(private val bazelInfo: BazelInfo) {
 
     fun resolveUri(path: Path): URI = uris.computeIfAbsent(path, Path::toUri)
 
+    fun unresolvedWorkspaceRoot(): Path = bazelInfo.workspaceRoot
+
     fun workspaceRoot(): URI = resolveUri(bazelInfo.workspaceRoot.toAbsolutePath())
 
     fun resolveUris(fileLocations: List<FileLocation>, shouldFilterExisting: Boolean = false): List<URI> =

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BspProjectMapper.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/BspProjectMapper.kt
@@ -17,6 +17,7 @@ import ch.epfl.scala.bsp4j.JavacOptionsItem
 import ch.epfl.scala.bsp4j.JavacOptionsParams
 import ch.epfl.scala.bsp4j.JavacOptionsResult
 import ch.epfl.scala.bsp4j.JvmEnvironmentItem
+import ch.epfl.scala.bsp4j.JvmMainClass
 import ch.epfl.scala.bsp4j.JvmRunEnvironmentParams
 import ch.epfl.scala.bsp4j.JvmRunEnvironmentResult
 import ch.epfl.scala.bsp4j.JvmTestEnvironmentParams
@@ -26,7 +27,6 @@ import ch.epfl.scala.bsp4j.OutputPathItemKind
 import ch.epfl.scala.bsp4j.OutputPathsItem
 import ch.epfl.scala.bsp4j.OutputPathsParams
 import ch.epfl.scala.bsp4j.OutputPathsResult
-import ch.epfl.scala.bsp4j.PythonBuildTarget
 import ch.epfl.scala.bsp4j.PythonOptionsItem
 import ch.epfl.scala.bsp4j.PythonOptionsParams
 import ch.epfl.scala.bsp4j.PythonOptionsResult
@@ -47,13 +47,16 @@ import ch.epfl.scala.bsp4j.SourcesParams
 import ch.epfl.scala.bsp4j.SourcesResult
 import ch.epfl.scala.bsp4j.TestProvider
 import ch.epfl.scala.bsp4j.WorkspaceBuildTargetsResult
+import org.eclipse.lsp4j.jsonrpc.CancelChecker
 import org.jetbrains.bsp.BazelBuildServerCapabilities
 import org.jetbrains.bsp.DirectoryItem
 import org.jetbrains.bsp.LibraryItem
 import org.jetbrains.bsp.WorkspaceDirectoriesResult
 import org.jetbrains.bsp.WorkspaceInvalidTargetsResult
 import org.jetbrains.bsp.WorkspaceLibrariesResult
+import org.jetbrains.bsp.bazel.bazelrunner.BazelRunner
 import org.jetbrains.bsp.bazel.commons.Constants
+import org.jetbrains.bsp.bazel.server.bsp.info.BspInfo
 import org.jetbrains.bsp.bazel.server.sync.languages.LanguagePluginsService
 import org.jetbrains.bsp.bazel.server.sync.languages.jvm.javaModule
 import org.jetbrains.bsp.bazel.server.sync.model.Label
@@ -64,12 +67,16 @@ import org.jetbrains.bsp.bazel.server.sync.model.Tag
 import org.jetbrains.bsp.bazel.workspacecontext.WorkspaceContextProvider
 import java.net.URI
 import java.nio.file.Path
+import java.nio.file.Paths
 import kotlin.io.path.name
 import kotlin.io.path.toPath
 
 class BspProjectMapper(
-    private val languagePluginsService: LanguagePluginsService,
-    private val workspaceContextProvider: WorkspaceContextProvider
+        private val languagePluginsService: LanguagePluginsService,
+        private val workspaceContextProvider: WorkspaceContextProvider,
+        private val bazelPathsResolver: BazelPathsResolver,
+        private val bazelRunner: BazelRunner,
+        private val bspInfo: BspInfo,
 ) {
 
     fun initializeServer(supportedLanguages: Set<Language>): InitializeBuildResult {
@@ -269,35 +276,62 @@ class BspProjectMapper(
     }
 
     fun jvmRunEnvironment(
-        project: Project, params: JvmRunEnvironmentParams
+            project: Project, params: JvmRunEnvironmentParams, cancelChecker: CancelChecker
     ): JvmRunEnvironmentResult {
         val targets = params.targets
-        val result = getJvmEnvironmentItems(project, targets)
+        val result = getJvmEnvironmentItems(project, targets, cancelChecker)
         return JvmRunEnvironmentResult(result)
     }
 
     fun jvmTestEnvironment(
-        project: Project, params: JvmTestEnvironmentParams
+            project: Project, params: JvmTestEnvironmentParams, cancelChecker: CancelChecker
     ): JvmTestEnvironmentResult {
         val targets = params.targets
-        val result = getJvmEnvironmentItems(project, targets)
+        val result = getJvmEnvironmentItems(project, targets, cancelChecker)
         return JvmTestEnvironmentResult(result)
     }
 
     private fun getJvmEnvironmentItems(
-        project: Project, targets: List<BuildTargetIdentifier>
+            project: Project, targets: List<BuildTargetIdentifier>, cancelChecker: CancelChecker
     ): List<JvmEnvironmentItem> {
-        fun extractJvmEnvironmentItem(module: Module): JvmEnvironmentItem? =
-            module.javaModule?.let {
-                languagePluginsService.javaLanguagePlugin.toJvmEnvironmentItem(module, it)
+        fun extractJvmEnvironmentItem(module: Module, runtimeClasspath: List<URI>): JvmEnvironmentItem? {
+            return module.javaModule?.let { javaModule ->
+                JvmEnvironmentItem(
+                        BspMappings.toBspId(module),
+                        runtimeClasspath.map { it.toString() },
+                        javaModule.jvmOps.toList(),
+                        bazelPathsResolver.unresolvedWorkspaceRoot().toString(),
+                        module.environmentVariables
+                ).apply {
+                    mainClasses = javaModule.mainClass?.let { listOf(JvmMainClass(it, javaModule.args)) }.orEmpty()
+                }
             }
+        }
 
-        val labels = BspMappings.toLabels(targets)
-        return labels.mapNotNull {
-            project.findModule(it)?.let(::extractJvmEnvironmentItem)
+        return targets.mapNotNull {
+            val label = Label(it.uri)
+            val module = project.findModule(label)
+            val runtimeClasspath = getRuntimeClasspath(cancelChecker, it)
+            module?.let { extractJvmEnvironmentItem(module, runtimeClasspath) }
         }
     }
 
+    private fun getRuntimeClasspath(cancelChecker: CancelChecker, target: BuildTargetIdentifier): List<URI> {
+        val queryFile = bspInfo.bazelBspDir().resolve("aspects/runtime_classpath_query.bzl")
+        val cqueryResult = bazelRunner.commandBuilder().cquery()
+                .withTargets(listOf(target.uri))
+                .withFlags(listOf("--starlark:file=$queryFile", "--output=starlark"))
+                .executeBazelCommand(parseProcessOutput = false)
+                .waitAndGetResult(cancelChecker, ensureAllOutputRead = true)
+        if (cqueryResult.isNotSuccess) throw RuntimeException("Could not query target '${target.uri}' for runtime classpath")
+        val runtimeClasspath = cqueryResult.stdoutLines
+                .filterNot { it.isEmpty() }
+                .map { bazelPathsResolver.resolveOutput(Paths.get(it))}
+                .filter { it.toFile().exists() } // I'm surprised this is needed, but we literally test it in e2e tests
+                .map { it.toUri() }
+        return runtimeClasspath
+    }
+
     fun buildTargetJavacOptions(project: Project, params: JavacOptionsParams): JavacOptionsResult {
         fun extractJavacOptionsItem(module: Module): JavacOptionsItem? =
             module.javaModule?.let {

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/ProjectSyncService.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/ProjectSyncService.kt
@@ -103,12 +103,12 @@ class ProjectSyncService(private val bspMapper: BspProjectMapper, private val pr
 
   fun jvmRunEnvironment(cancelChecker: CancelChecker, params: JvmRunEnvironmentParams): JvmRunEnvironmentResult {
     val project = projectProvider.get(cancelChecker)
-    return bspMapper.jvmRunEnvironment(project, params)
+    return bspMapper.jvmRunEnvironment(project, params, cancelChecker)
   }
 
   fun jvmTestEnvironment(cancelChecker: CancelChecker, params: JvmTestEnvironmentParams): JvmTestEnvironmentResult {
     val project = projectProvider.get(cancelChecker)
-    return bspMapper.jvmTestEnvironment(project, params)
+    return bspMapper.jvmTestEnvironment(project, params, cancelChecker)
   }
 
   fun buildTargetJavacOptions(cancelChecker: CancelChecker, params: JavacOptionsParams): JavacOptionsResult {

diff --git a/.../src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaLanguagePlugin.kt b/.../src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaLanguagePlugin.kt
@@ -57,7 +57,6 @@ class JavaLanguagePlugin(
                 allOutputs,
                 mainClass,
                 argsList,
-                runtimeClasspath,
                 compileClasspath,
                 sourcesClasspath,
                 ideClasspath
@@ -112,17 +111,6 @@ class JavaLanguagePlugin(
         }
     }
 
-    fun toJvmEnvironmentItem(module: Module, javaModule: JavaModule): JvmEnvironmentItem =
-        JvmEnvironmentItem(
-            BspMappings.toBspId(module),
-            javaModule.runtimeClasspath.map { it.toString() }.toList(),
-            javaModule.jvmOps.toList(),
-            bazelInfo.workspaceRoot.toString(),
-            module.environmentVariables
-        ).apply {
-            mainClasses = javaModule.mainClass?.let { listOf(JvmMainClass(it, javaModule.args)) }.orEmpty()
-        }
-
     fun toJavacOptionsItem(module: Module, javaModule: JavaModule): JavacOptionsItem =
         JavacOptionsItem(
             BspMappings.toBspId(module),

diff --git a/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaModule.kt b/server/src/main/kotlin/org/jetbrains/bsp/bazel/server/sync/languages/java/JavaModule.kt
@@ -18,7 +18,6 @@ data class JavaModule(
     val allOutputs: List<URI>,
     val mainClass: String?,
     val args: List<String>,
-    val runtimeClasspath: List<URI>,
     val compileClasspath: List<URI>,
     val sourcesClasspath: List<URI>,
     val ideClasspath: List<URI>

diff --git a/server/src/test/kotlin/org/jetbrains/bsp/bazel/server/sync/ProjectStorageTest.kt b/server/src/test/kotlin/org/jetbrains/bsp/bazel/server/sync/ProjectStorageTest.kt
@@ -40,7 +40,6 @@ class ProjectStorageTest {
                 emptyList(),
                 emptyList(),
                 emptyList(),
-                emptyList(),
                 emptyList()
             )
         )