improvement: Request jvmEnvironment lazily

Previously, we would eagerly request jvmEnvironment for all known main classes, which in case of Bazel might cause a lot of queries. This might also not be efficient in some other cases. Now, we only request and cache jvmRunEnvironment when needed.
scalameta · Mar 18, 2024 · fe1aff2 · fe1aff2
1 parent b8e1fc1
commit fe1aff2
Show file tree

Hide file tree

Showing 10 changed files with 119 additions and 44 deletions.
diff --git a/metals/src/main/scala/scala/meta/internal/metals/CodeLensProvider.scala b/metals/src/main/scala/scala/meta/internal/metals/CodeLensProvider.scala
@@ -1,5 +1,8 @@
 package scala.meta.internal.metals
 
+import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
+
 import scala.meta.internal.implementation.TextDocumentWithPath
 import scala.meta.internal.metals.codelenses.CodeLens
 import scala.meta.internal.mtags.Semanticdbs
@@ -11,23 +14,23 @@ final class CodeLensProvider(
     codeLensProviders: List[CodeLens],
     semanticdbs: Semanticdbs,
     stacktraceAnalyzer: StacktraceAnalyzer,
-) {
+)(implicit val ec: ExecutionContext) {
   // code lenses will be refreshed after compilation or when workspace gets indexed
-  def findLenses(path: AbsolutePath): Seq[l.CodeLens] = {
+  def findLenses(path: AbsolutePath): Future[Seq[l.CodeLens]] = {
     if (stacktraceAnalyzer.isStackTraceFile(path)) {
-      stacktraceAnalyzer.stacktraceLenses(path)
+      Future(stacktraceAnalyzer.stacktraceLenses(path))
     } else {
       val enabledCodelenses = codeLensProviders.filter(_.isEnabled)
       val semanticdbCodeLenses = semanticdbs
         .textDocument(path)
         .documentIncludingStale
         .map { textDocument =>
           val doc = TextDocumentWithPath(textDocument, path)
-          enabledCodelenses.flatMap(_.codeLenses(doc))
+          enabledCodelenses.map(_.codeLenses(doc))
         }
         .getOrElse(Seq.empty)
-      val otherCodeLenses = enabledCodelenses.flatMap(_.codeLenses(path))
-      semanticdbCodeLenses ++ otherCodeLenses
+      val otherCodeLenses = enabledCodelenses.map(_.codeLenses(path))
+      Future.sequence(semanticdbCodeLenses ++ otherCodeLenses).map(_.flatten)
     }
   }
 }
diff --git a/metals/src/main/scala/scala/meta/internal/metals/Indexer.scala b/metals/src/main/scala/scala/meta/internal/metals/Indexer.scala
@@ -216,6 +216,7 @@ final case class Indexer(
         val data = buildTool.data
         val importedBuild = buildTool.importedBuild
         data.reset()
+        buildTargetClasses.clear()
         data.addWorkspaceBuildTargets(importedBuild.workspaceBuildTargets)
         data.addScalacOptions(
           importedBuild.scalacOptions,

diff --git a/metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala b/metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala
@@ -1631,13 +1631,13 @@ class MetalsLspService(
       params: CodeLensParams
   ): CompletableFuture[util.List[CodeLens]] =
     CancelTokens.future { _ =>
-      buildServerPromise.future.map { _ =>
+      buildServerPromise.future.flatMap { _ =>
         timerProvider.timedThunk(
           "code lens generation",
           thresholdMillis = 1.second.toMillis,
         ) {
           val path = params.getTextDocument.getUri.toAbsolutePath
-          codeLensProvider.findLenses(path).toList.asJava
+          codeLensProvider.findLenses(path).map(_.toList.asJava)
         }
       }
     }

diff --git a/metals/src/main/scala/scala/meta/internal/metals/codelenses/CodeLens.scala b/metals/src/main/scala/scala/meta/internal/metals/codelenses/CodeLens.scala
@@ -1,5 +1,7 @@
 package scala.meta.internal.metals.codelenses
 
+import scala.concurrent.Future
+
 import scala.meta.internal.implementation.TextDocumentWithPath
 import scala.meta.io.AbsolutePath
 
@@ -8,8 +10,11 @@ import org.eclipse.{lsp4j => l}
 trait CodeLens {
   def isEnabled: Boolean
 
-  def codeLenses(textDocumentWithPath: TextDocumentWithPath): Seq[l.CodeLens] =
-    Seq.empty
+  def codeLenses(
+      textDocumentWithPath: TextDocumentWithPath
+  ): Future[Seq[l.CodeLens]] =
+    Future.successful(Seq.empty)
 
-  def codeLenses(path: AbsolutePath): Seq[l.CodeLens] = Seq.empty
+  def codeLenses(path: AbsolutePath): Future[Seq[l.CodeLens]] =
+    Future.successful(Seq.empty)
 }
diff --git a/metals/src/main/scala/scala/meta/internal/metals/codelenses/RunTestCodeLens.scala b/metals/src/main/scala/scala/meta/internal/metals/codelenses/RunTestCodeLens.scala
@@ -2,6 +2,9 @@ package scala.meta.internal.metals.codelenses
 
 import java.util.Collections.singletonList
 
+import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
+
 import scala.meta.internal.implementation.TextDocumentWithPath
 import scala.meta.internal.metals.BaseCommand
 import scala.meta.internal.metals.Buffers
@@ -50,14 +53,31 @@ final class RunTestCodeLens(
     userConfig: () => UserConfiguration,
     trees: Trees,
     workspace: AbsolutePath,
-) extends CodeLens {
+)(implicit val ec: ExecutionContext)
+    extends CodeLens {
 
   override def isEnabled: Boolean =
     clientConfig.isDebuggingProvider() || clientConfig.isRunProvider()
 
   override def codeLenses(
       textDocumentWithPath: TextDocumentWithPath
-  ): Seq[l.CodeLens] = {
+  ): Future[Seq[l.CodeLens]] = {
+
+    /**
+     * Jvm environment needs to be requested lazily.
+     * This requests and caches it for later use, otherwise we
+     * would need to forward Future across different methods
+     * which would make things way too complex.
+     */
+    def requestJvmEnvironment(
+        buildTargetId: BuildTargetIdentifier,
+        isJvm: Boolean,
+    ): Future[Unit] = {
+      if (isJvm)
+        buildTargetClasses.jvmRunEnvironment(buildTargetId).map(_ => ())
+      else
+        Future.unit
+    }
     val textDocument = textDocumentWithPath.textDocument
     val path = textDocumentWithPath.filePath
     val distance = buffers.tokenEditDistance(path, textDocument.text, trees)
@@ -71,8 +91,9 @@ final class RunTestCodeLens(
       // although hasDebug is already available in BSP capabilities
       // see https://github.com/build-server-protocol/build-server-protocol/pull/161
       // most of the bsp servers such as bloop and sbt might not support it.
-    } yield {
+    } yield requestJvmEnvironment(buildTargetId, isJVM).map { _ =>
       val classes = buildTargetClasses.classesOf(buildTargetId)
+
       // sbt doesn't declare debugging provider
       def buildServerCanDebug =
         connection.isDebuggingProvider || connection.isSbt
@@ -100,7 +121,7 @@ final class RunTestCodeLens(
 
     }
 
-    lenses.getOrElse(Seq.empty)
+    lenses.getOrElse(Future.successful(Nil))
   }
 
   /**
@@ -315,8 +336,8 @@ final class RunTestCodeLens(
     val (data, shellCommandAdded) =
       if (!isJVM) (main.toJson, false)
       else
-        buildTargetClasses.jvmRunEnvironment
-          .get(target)
+        buildTargetClasses
+          .jvmRunEnvironmentSync(target)
           .zip(javaBinary) match {
           case None =>
             (main.toJson, false)

diff --git a/metals/src/main/scala/scala/meta/internal/metals/codelenses/SuperMethodCodeLens.scala b/metals/src/main/scala/scala/meta/internal/metals/codelenses/SuperMethodCodeLens.scala
@@ -1,5 +1,8 @@
 package scala.meta.internal.metals.codelenses
 
+import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
+
 import scala.meta.internal.implementation.SuperMethodProvider
 import scala.meta.internal.implementation.TextDocumentWithPath
 import scala.meta.internal.metals.Buffers
@@ -21,13 +24,14 @@ final class SuperMethodCodeLens(
     userConfig: () => UserConfiguration,
     clientConfig: ClientConfiguration,
     trees: Trees,
-) extends CodeLens {
+)(implicit val ec: ExecutionContext)
+    extends CodeLens {
 
   override def isEnabled: Boolean = userConfig().superMethodLensesEnabled
 
   override def codeLenses(
       textDocumentWithPath: TextDocumentWithPath
-  ): Seq[l.CodeLens] = {
+  ): Future[Seq[l.CodeLens]] = Future {
     val textDocument = textDocumentWithPath.textDocument
     val path = textDocumentWithPath.filePath
 

diff --git a/metals/src/main/scala/scala/meta/internal/metals/codelenses/WorksheetCodeLens.scala b/metals/src/main/scala/scala/meta/internal/metals/codelenses/WorksheetCodeLens.scala
@@ -1,19 +1,24 @@
 package scala.meta.internal.metals.codelenses
 
+import scala.concurrent.ExecutionContext
+import scala.concurrent.Future
+
 import scala.meta.internal.metals.ClientCommands.CopyWorksheetOutput
 import scala.meta.internal.metals.ClientConfiguration
 import scala.meta.internal.metals.MetalsEnrichments._
 import scala.meta.io.AbsolutePath
 
 import org.eclipse.{lsp4j => l}
 
-class WorksheetCodeLens(clientConfig: ClientConfiguration) extends CodeLens {
+class WorksheetCodeLens(clientConfig: ClientConfiguration)(implicit
+    val ec: ExecutionContext
+) extends CodeLens {
 
   override def isEnabled: Boolean = clientConfig.isCopyWorksheetOutputProvider
 
   override def codeLenses(
       path: AbsolutePath
-  ): Seq[l.CodeLens] = {
+  ): Future[Seq[l.CodeLens]] = Future {
     if (path.isWorksheet) {
       val command = CopyWorksheetOutput.toLsp(path.toURI)
       val startPosition = new l.Position(0, 0)

diff --git a/metals/src/main/scala/scala/meta/internal/metals/debug/BuildTargetClasses.scala b/metals/src/main/scala/scala/meta/internal/metals/debug/BuildTargetClasses.scala
@@ -20,7 +20,7 @@ final class BuildTargetClasses(
     buildTargets: BuildTargets
 )(implicit val ec: ExecutionContext) {
   private val index = TrieMap.empty[b.BuildTargetIdentifier, Classes]
-  val jvmRunEnvironment
+  private val jvmRunEnvironments
       : TrieMap[b.BuildTargetIdentifier, b.JvmEnvironmentItem] =
     TrieMap.empty[b.BuildTargetIdentifier, b.JvmEnvironmentItem]
   val rebuildIndex: BatchedFunction[b.BuildTargetIdentifier, Unit] =
@@ -38,6 +38,10 @@ final class BuildTargetClasses(
     index.put(target, new Classes)
   }
 
+  def clear(): Unit = {
+    jvmRunEnvironments.clear()
+  }
+
   def findMainClassByName(
       name: String
   ): List[(b.ScalaMainClass, b.BuildTargetIdentifier)] =
@@ -88,14 +92,9 @@ final class BuildTargetClasses(
                 .map(cacheTestClasses(classes, _))
             else Future.unit
 
-          val jvmRunEnvironment = connection
-            .jvmRunEnvironment(new b.JvmRunEnvironmentParams(targetsList))
-            .map(cacheJvmRunEnvironment)
-
           for {
             _ <- updateMainClasses
             _ <- updateTestClasses
-            _ <- jvmRunEnvironment
           } yield {
             classes.foreach { case (id, classes) =>
               index.put(id, classes)
@@ -105,14 +104,41 @@ final class BuildTargetClasses(
       .ignoreValue
   }
 
+  def jvmRunEnvironmentSync(
+      buildTargetId: b.BuildTargetIdentifier
+  ): Option[b.JvmEnvironmentItem] = jvmRunEnvironments.get(buildTargetId)
+
+  def jvmRunEnvironment(
+      buildTargetId: b.BuildTargetIdentifier
+  ): Future[Option[b.JvmEnvironmentItem]] = {
+    jvmRunEnvironments.get(buildTargetId) match {
+      case None =>
+        buildTargets.buildServerOf(buildTargetId) match {
+          case None => Future.successful(None)
+          case Some(connection) =>
+            connection
+              .jvmRunEnvironment(
+                new b.JvmRunEnvironmentParams(List(buildTargetId).asJava)
+              )
+              .map { env =>
+                cacheJvmRunEnvironment(env)
+                env.getItems().asScala.headOption
+              }
+
+        }
+      case jvmRunEnv: Some[b.JvmEnvironmentItem] =>
+        Future.successful(jvmRunEnv)
+    }
+  }
+
   private def cacheJvmRunEnvironment(
       result: b.JvmRunEnvironmentResult
   ): Unit = {
     for {
       item <- result.getItems().asScala
       target = item.getTarget
     } {
-      jvmRunEnvironment.put(target, item)
+      jvmRunEnvironments.put(target, item)
     }
   }
 

diff --git a/metals/src/main/scala/scala/meta/internal/metals/debug/DebugProvider.scala b/metals/src/main/scala/scala/meta/internal/metals/debug/DebugProvider.scala
@@ -519,13 +519,13 @@ class DebugProvider(
   def runCommandDiscovery(
       unresolvedParams: DebugDiscoveryParams
   )(implicit ec: ExecutionContext): Future[b.DebugSessionParams] = {
-    debugDiscovery(unresolvedParams).map(enrichWithMainShellCommand)
+    debugDiscovery(unresolvedParams).flatMap(enrichWithMainShellCommand)
   }
 
   private def enrichWithMainShellCommand(
       params: b.DebugSessionParams
-  ): b.DebugSessionParams = {
-    params.getData() match {
+  )(implicit ec: ExecutionContext): Future[b.DebugSessionParams] = {
+    val future = params.getData() match {
       case json: JsonElement
           if params.getDataKind == b.DebugSessionParamsDataKind.SCALA_MAIN_CLASS =>
         json.as[b.ScalaMainClass] match {
@@ -536,21 +536,31 @@ class DebugProvider(
                 JavaBinary.javaBinaryFromPath(scalaTarget.jvmHome)
               )
               .orElse(userConfig().usedJavaBinary)
-            val updatedData = buildTargetClasses.jvmRunEnvironment
-              .get(params.getTargets().get(0))
-              .zip(javaBinary) match {
-              case None =>
-                main.toJson
-              case Some((env, javaHome)) =>
-                ExtendedScalaMainClass(main, env, javaHome, workspace).toJson
-            }
-            params.setData(updatedData)
-          case _ =>
+            buildTargetClasses
+              .jvmRunEnvironment(params.getTargets().get(0))
+              .map { envItem =>
+                val updatedData = envItem.zip(javaBinary) match {
+                  case None =>
+                    main.toJson
+                  case Some((env, javaHome)) =>
+                    ExtendedScalaMainClass(
+                      main,
+                      env,
+                      javaHome,
+                      workspace,
+                    ).toJson
+                }
+                params.setData(updatedData)
+              }
+          case _ => Future.unit
         }
 
-      case _ =>
+      case _ => Future.unit
+    }
+
+    future.map { _ =>
+      params
     }
-    params
   }
 
   /**

diff --git a/metals/src/main/scala/scala/meta/internal/metals/testProvider/TestSuitesProvider.scala b/metals/src/main/scala/scala/meta/internal/metals/testProvider/TestSuitesProvider.scala
@@ -130,7 +130,7 @@ final class TestSuitesProvider(
 
   override def codeLenses(
       textDocumentWithPath: TextDocumentWithPath
-  ): Seq[l.CodeLens] = {
+  ): Future[Seq[l.CodeLens]] = Future {
     val path = textDocumentWithPath.filePath
     for {
       target <- buildTargets.inverseSources(path).toList