Fix stale Java diagnostics after MBT import

metals/src/main/scala/scala/meta/internal/metals/Compilers.scala

@@ -7,6 +7,7 @@ import java.time.Duration
 import java.util.concurrent.ScheduledExecutorService
 import java.util.concurrent.ScheduledFuture
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicLong
 import java.{util => ju}
 
 import scala.annotation.nowarn
@@ -313,6 +314,7 @@ class Compilers(
   override def cancel(): Unit = {
     presentationCompilerCache.invalidateAll()
     presentationCompilerWorksheetsCache.invalidateAll()
+    presentationCompilerGeneration.incrementAndGet()
     worksheetsDigests.clear()
     outlineFilesProvider.clear()
   }
@@ -353,7 +355,11 @@ class Compilers(
     loadCompiler(path).foreach(_.didClose(path.toNIO.toUri()))
   }
 
-  def didFocus(path: AbsolutePath): Future[List[Diagnostic]] = {
+  def didFocus(
+      path: AbsolutePath,
+      retryStale: Boolean = true,
+  ): Future[List[Diagnostic]] = {
+    val generation = presentationCompilerGeneration.get()
     val maybeDiagnostics =
       for (pc <- loadCompiler(path); contents <- buffers.get(path))
         yield {
@@ -385,24 +391,37 @@ class Compilers(
               )
               result
             }
+            .flatMap { result =>
+              if (generation == presentationCompilerGeneration.get())
+                Future.successful(result)
+              else if (retryStale) didFocus(path, retryStale = false)
+              else Future.successful(Nil)
+            }
         }
 
     maybeDiagnostics.getOrElse(Future.successful(List.empty))
   }
 
+  private val presentationCompilerGeneration = new AtomicLong(0)
   private val inFlightDidChange =
     TrieMap.empty[AbsolutePath, CompletableCancelToken]
+  private case class DidChangeRequest(
+      path: AbsolutePath,
+      generation: Long,
+  )
   private val diagnosticsDebouncerDelay: FiniteDuration =
     if (Testing.isEnabled) 0.millis
     else sys.Prop[Int]("metals.errors-delay").option.getOrElse(500).millis
-  private val fileDidChange: BatchedFunction[AbsolutePath, Unit] =
-    BatchedFunction.fromFuture[AbsolutePath, Unit](
-      changedFiles => {
+  private val fileDidChange: BatchedFunction[DidChangeRequest, Unit] =
+    BatchedFunction.fromFuture[DidChangeRequest, Unit](
+      requests => {
         for {
           _ <- sh.sleep(diagnosticsDebouncerDelay)
+          files = requests
+            .groupMapReduce(_.path)(_.generation)(math.max)
           futures =
             for {
-              file <- changedFiles.distinct
+              (file, generation) <- files
               pc <- this.loadCompiler(file).toList
               contents <- buffers.get(file).toList
             } yield {
@@ -426,15 +445,20 @@ class Compilers(
                   pc.didChange(params).asScala
                 }
                 .map { case (timer, reportedDiagnostics) =>
-                  diagnostics.publishDiagnosticsNotAdjusted(
-                    file,
-                    reportedDiagnostics.asScala.toList,
-                  )
-                  metrics.recordEvent(
-                    Event
-                      .duration("diagnostics", timer.elapsed)
-                      .withLanguage(file.toJLanguage)
-                  )
+                  if (generation == presentationCompilerGeneration.get()) {
+                    diagnostics.publishDiagnosticsNotAdjusted(
+                      file,
+                      reportedDiagnostics.asScala.toList,
+                    )
+                    metrics.recordEvent(
+                      Event
+                        .duration("diagnostics", timer.elapsed)
+                        .withLanguage(file.toJLanguage)
+                    )
+                  }
+                }
+                .andThen { case _ =>
+                  inFlightDidChange.remove(file, token)
                 }
             }
           _ <- Future.sequence(futures)
@@ -553,7 +577,9 @@ class Compilers(
   def didChange(path: AbsolutePath): Future[Unit] = {
     if (userConfig().presentationCompilerDiagnostics)
       // Batch/debounce these requests since they can arrive in bursts
-      fileDidChange(Seq(path))
+      fileDidChange(
+        Seq(DidChangeRequest(path, presentationCompilerGeneration.get()))
+      )
     else
       didChangeBSPDiagnostics(path, shouldReturnDiagnostics = false).ignoreValue
   }
@@ -653,7 +679,11 @@ class Compilers(
     // Restart PC for all build targets that depend on this target
     for {
       target <- buildTargets.allInverseDependencies(target)
-      compiler <- buildTargetPCFromCache(target)
+      key <- List(
+        PresentationCompilerKey.ScalaBuildTarget(target),
+        PresentationCompilerKey.JavaBuildTarget(target),
+      )
+      compiler <- cache.get(key).map(_.await)
     } {
       scribe.debug(s"Restarting PC for target ${target.getUri}")
       compiler.restart()

metals/src/main/scala/scala/meta/internal/metals/ConnectionProvider.scala

@@ -70,6 +70,7 @@ class ConnectionProvider(
     indexProviders: IndexProviders,
     syncStatusReporter: SyncStatusReporter,
     mbtBuild: () => MbtBuild,
+    refreshMbtStateAfterIndex: () => Future[Unit],
     mbtDebugStarter: () => Option[MbtDebugSessionStarter] = () => None,
 )(implicit ec: ExecutionContextExecutorService, rc: ReportContext)
     extends Indexer(indexProviders, mbtBuild)
@@ -253,17 +254,22 @@ class ConnectionProvider(
       case Some(session) if session.canReloadWorkspace =>
         workDoneProgress.trackProgressFuture(
           "Sync",
-          progress =>
-            for {
-              _ <- session.workspaceReload()
-              _ <- connect(new ImportBuildAndIndex(session), progress)
-            } yield (),
+          progress => reloadAndImport(session, progress).ignoreValue,
           metricName = Some("reload_build_server"),
         )
       case _ =>
         fullConnect()
     }
 
+  private def reloadAndImport(
+      session: BspSession,
+      progress: TaskProgress,
+  ): Future[BuildChange] =
+    for {
+      _ <- session.workspaceReload()
+      buildChange <- connect(new ImportBuildAndIndex(session), progress)
+    } yield buildChange
+
   private def isBspAvailable(buildTool: BuildTool) =
     buildTool.isBspGenerated(folder) || bspGlobalDirectories.exists(
       _.resolve(s"${buildTool.buildServerName}.json").isFile
@@ -285,6 +291,10 @@ class ConnectionProvider(
         else runMbtReimport(mbtImporters)
       runImport.flatMap { _ =>
         bspSession match {
+          case Some(session)
+              if MbtBuildServer.isMbtServer(session.main.name) &&
+                session.canReloadWorkspace =>
+            reloadAndImport(session, progress)
           case Some(session) =>
             connect(new ImportBuildAndIndex(session), progress)
           case None =>
@@ -680,32 +690,27 @@ class ConnectionProvider(
         }
         _ = compilers.cancel()
         buildChange <- index(check, progress)
-        // When testing we need to make sure the classpath is refreshed after mbt.json is generated
-        _ <- {
-          if (MetalsServerConfig.isTesting)
-            refreshMbtTurbineClasspath(session).withInterrupt
-          else
-            Future {
-              refreshMbtTurbineClasspath(session)
-            }.withInterrupt
-        }
+        _ <- refreshMbtStateAfterIndexIfNeeded(session)
       } yield {
         syncStatusReporter.importFinished(focusedDocument.map(_.toURI.toString))
         buildChange
       }
     }
 
-    private def refreshMbtTurbineClasspath(
+    private def refreshMbtStateAfterIndexIfNeeded(
         session: BspSession
-    ): Future[Unit] =
-      if (
-        MbtBuildServer.isMbtServer(session.main.name) &&
-        userConfig.javaSymbolLoader.isTurbineClasspath
-      ) {
-        mbtSymbolSearch.scheduleRecompileTurbineClasspath()
-      } else {
-        Future.unit
-      }
+    ): Interruptable[Unit] = {
+      val refresh =
+        if (MbtBuildServer.isMbtServer(session.main.name))
+          refreshMbtStateAfterIndex()
+        else Future.unit
+
+      if (MetalsServerConfig.isTesting) refresh.withInterrupt
+      else
+        refresh.recover { case error =>
+          scribe.warn("failed to refresh MBT diagnostics", error)
+        }.withInterrupt
+    }
 
     private def saveProjectReferencesInfo(
         bspBuilds: List[BspSession.BspBuild]

metals/src/main/scala/scala/meta/internal/metals/MetalsLspService.scala

@@ -1157,20 +1157,31 @@ abstract class MetalsLspService(
   private def refreshAllDiagnostics(): Future[Unit] = {
     refreshDiagnostics(_ => true)
   }
+
+  protected def refreshMbtStateAfterIndex(): Future[Unit] =
+    for {
+      _ <-
+        if (userConfig.javaSymbolLoader.isTurbineClasspath)
+          mbt2.recompileTurbineClasspath()
+        else Future.unit
+      _ = compilers.cancel()
+      _ = diagnostics.reset(buffers.open.toSeq)
+      _ <- refreshAllDiagnostics()
+    } yield ()
+
   protected def refreshDiagnostics(
       isIncludedPath: AbsolutePath => Boolean
-  ): Future[Unit] = {
+  ): Future[Unit] =
     // rerun diagnostics for all open documents
-    val futures =
-      buffers.open.filter(isIncludedPath).map { path =>
+    Future
+      .traverse(buffers.open.filter(isIncludedPath)) { path =>
         for {
           reportedDiagnostics <- compilers.didFocus(path)
           _ = diagnostics
             .publishDiagnosticsNotAdjusted(path, reportedDiagnostics)
         } yield ()
       }
-    Future.sequence(futures).map(_ => ())
-  }
+      .ignoreValue
 
   def resetPresentationCompilers(): Future[Unit] = {
     compilers.restartAll()

metals/src/main/scala/scala/meta/internal/metals/ProjectMetalsLspService.scala

@@ -237,6 +237,7 @@ class ProjectMetalsLspService(
       this,
       syncStatusReporter,
       () => mbtBuild,
+      () => refreshMbtStateAfterIndex(),
       mbtDebugStarter = () => mbtDebugStarter,
     )
     provider.buildServerPromise.future.onComplete(_ => moduleStatus.refresh())

metals/src/main/scala/scala/meta/internal/metals/mbt/TurbineCompiler.scala

@@ -5,6 +5,7 @@ import java.nio.file.Path
 import java.util.Optional
 import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.atomic.AtomicBoolean
+import java.util.concurrent.atomic.AtomicLong
 import java.{util => ju}
 import javax.tools.JavaFileManager
 import javax.tools.JavaFileObject
@@ -156,6 +157,8 @@ class TurbineCompiler[T](
   // In this mode, we rely entirely on SOURCE_PATH fallback for updated sources.
   private def isRecompilationDisabled: Boolean =
     debounceDelay.toMillis >= 3600000
+  private val compileGeneration = new AtomicLong(0L)
+  private val compileLock = new Object
 
   private val doCompile =
     BatchedFunction.fromFuture[Unit, TurbineCompileResult](
@@ -166,12 +169,14 @@ class TurbineCompiler[T](
           Future.successful(result)
         } else {
           val toCompile = sourcepathSources()
+          val generation = compileGeneration.get()
           for {
             _ <- sleeper.sleep(debounceDelay)
           } yield {
-            val result = doCompileNow()
-            toCompile.foreach(_.isCompiled.set(true))
-            result
+            doCompileNow(
+              expectedGeneration = generation,
+              markCompiled = toCompile,
+            )
           }
         }
       },
@@ -183,20 +188,31 @@ class TurbineCompiler[T](
     )
   }
 
-  var result = TurbineCompiler.emptyResult
-  def doCompileNow(): TurbineCompileResult = {
-    result = TurbineCompiler.compileClassfiles(
-      allCompilationUnits(),
-      parseUnit,
-      classpath(),
-      progressBars,
-    )
-    cleanup()
-    // Clear deleted binary names after recompile - they are no longer in the compiled output
-    deletedBinaryNames.clear()
-    onIndexingDone()
-    result
-  }
+  @volatile var result = TurbineCompiler.emptyResult
+  private def doCompileNow(
+      expectedGeneration: Long,
+      markCompiled: Seq[SourcepathJavaFileObject],
+  ): TurbineCompileResult =
+    compileLock.synchronized {
+      if (expectedGeneration != compileGeneration.get()) result
+      else {
+        val compiled = TurbineCompiler.compileClassfiles(
+          allCompilationUnits(),
+          parseUnit,
+          classpath(),
+          progressBars,
+        )
+        if (expectedGeneration == compileGeneration.get()) {
+          result = compiled
+          markCompiled.foreach(_.isCompiled.set(true))
+          cleanup()
+          // Clear deleted binary names after recompile - they are no longer in the compiled output
+          deletedBinaryNames.clear()
+          onIndexingDone()
+        }
+        result
+      }
+    }
 
   /**
    * Called when a file is deleted. Tracks the binary names of the deleted classes
@@ -225,8 +241,16 @@ class TurbineCompiler[T](
     deletedBinaryNames.contains(binaryName)
   }
 
-  def compileNow(): Future[TurbineCompileResult] = Future {
-    doCompileNow()
+  def compileNow(): Future[TurbineCompileResult] = {
+    val toCompile = sourcepathSources()
+    val generation = compileGeneration.incrementAndGet()
+    doCompile.cancelAll()
+    Future {
+      doCompileNow(
+        expectedGeneration = generation,
+        markCompiled = toCompile,
+      )
+    }
   }
 
   def scheduleCompile(): Future[TurbineCompileResult] = {