Fix deadlock with inspect code

src/main/kotlin/build/buf/intellij/annotator/BufAnalyzeUtils.kt

@@ -19,7 +19,7 @@ import build.buf.intellij.config.BufConfig
 import build.buf.intellij.model.BufIssue
 import build.buf.intellij.settings.BufCLIUtils
 import build.buf.intellij.settings.bufSettings
-import build.buf.intellij.vendor.isProtobufFile
+import build.buf.intellij.vendor.protobufLanguage
 import com.intellij.execution.configurations.GeneralCommandLine
 import com.intellij.execution.process.OSProcessHandler
 import com.intellij.execution.process.ProcessEvent
@@ -36,11 +36,13 @@ import com.intellij.openapi.components.service
 import com.intellij.openapi.diagnostic.Logger
 import com.intellij.openapi.diagnostic.logger
 import com.intellij.openapi.diagnostic.thisLogger
+import com.intellij.openapi.editor.Document
 import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.openapi.fileTypes.LanguageFileType
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.Disposer
-import com.intellij.psi.PsiDocumentManager
+import com.intellij.openapi.vfs.VirtualFile
 import com.intellij.psi.util.PsiModificationTracker
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
@@ -88,10 +90,15 @@ object BufAnalyzeUtils {
         workingDirectory: Path,
     ): BufAnalyzeResult {
         ProgressManager.checkCanceled()
-        WriteAction.computeAndWait<Unit, Throwable> {
-            FileDocumentManager.getInstance().saveDocuments { document ->
-                val psiFile = PsiDocumentManager.getInstance(project).getPsiFile(document) ?: return@saveDocuments false
-                psiFile.isProtobufFile() || BufConfig.CONFIG_FILES.contains(psiFile.name)
+        // Saving needs a write action, so callers must not hold a read lock or
+        // wait on this from a thread that does. Skipping the write action when
+        // nothing relevant is unsaved keeps headless runs, which may hold a
+        // global read lock, from deadlocking.
+        val fileDocumentManager = FileDocumentManager.getInstance()
+        val isBufDocument = { document: Document -> fileDocumentManager.getFile(document)?.isBufFile() == true }
+        if (fileDocumentManager.unsavedDocuments.any(isBufDocument)) {
+            WriteAction.computeAndWait<Unit, Throwable> {
+                fileDocumentManager.saveDocuments(isBufDocument)
             }
         }
         val started = Instant.now()
@@ -141,6 +148,8 @@ object BufAnalyzeUtils {
         return BufAnalyzeResult(workingDirectory, issues)
     }
 
+    private fun VirtualFile.isBufFile(): Boolean = (fileType as? LanguageFileType)?.language == protobufLanguage() || name in BufConfig.CONFIG_FILES
+
     private fun findBreakingArguments(workingDirectory: Path, gitRepoRoot: Path): List<String> {
         if (gitRepoRoot == workingDirectory) {
             return listOf("--against", ".git")

src/main/kotlin/build/buf/intellij/inspections/BufAnalyzeInspection.kt

@@ -21,6 +21,7 @@ import build.buf.intellij.annotator.BufAnalyzeUtils
 import build.buf.intellij.annotator.addHighlightsForFile
 import build.buf.intellij.annotator.createDisposableOnAnyPsiChange
 import build.buf.intellij.vendor.isProtobufFile
+import build.buf.intellij.vendor.protobufLanguage
 import com.intellij.codeInsight.daemon.impl.HighlightInfo
 import com.intellij.codeInspection.GlobalInspectionContext
 import com.intellij.codeInspection.GlobalSimpleInspectionTool
@@ -31,23 +32,33 @@ import com.intellij.codeInspection.ProblemsHolder
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.application.runReadAction
 import com.intellij.openapi.components.service
+import com.intellij.openapi.progress.ProgressManager
+import com.intellij.openapi.progress.util.ProgressIndicatorUtils
 import com.intellij.openapi.roots.ProjectFileIndex
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.Key
 import com.intellij.psi.PsiFile
+import com.intellij.psi.util.PsiModificationTracker
+import java.util.concurrent.ConcurrentHashMap
 
-class BufAnalyzeInspection : GlobalSimpleInspectionTool() {
-    companion object {
-        private const val SHORT_NAME = "BufAnalyze"
-    }
+private val ANALYZED_FILES: Key<MutableSet<PsiFile>> = Key.create("BUF_ANALYZED_FILES")
+private const val MAX_ANALYZE_ATTEMPTS = 3
 
+class BufAnalyzeInspection : GlobalSimpleInspectionTool() {
     private val appService = service<BufPluginService>()
 
-    override fun getShortName(): String = SHORT_NAME
-
     override fun getDisplayName(): String = BufBundle.message("buf.inspection.analyze.display.name")
 
-    override fun isReadActionNeeded(): Boolean = false
+    override fun inspectionStarted(
+        manager: InspectionManager,
+        globalContext: GlobalInspectionContext,
+        problemDescriptionsProcessor: ProblemDescriptionsProcessor,
+    ) {
+        globalContext.putUserData(ANALYZED_FILES, ConcurrentHashMap.newKeySet())
+    }
 
+    // The platform calls checkFile inside a read action. Waiting on buf here
+    // deadlocks, because the analysis saves documents in a write action first.
     override fun checkFile(
         file: PsiFile,
         manager: InspectionManager,
@@ -56,21 +67,61 @@ class BufAnalyzeInspection : GlobalSimpleInspectionTool() {
         problemDescriptionsProcessor: ProblemDescriptionsProcessor,
     ) {
         if (!file.isProtobufFile()) return
-        val project = manager.project
-        val disposable = project.messageBus.createDisposableOnAnyPsiChange()
-            .also { Disposer.register(appService, it) }
-
-        val contentRootForFile = ProjectFileIndex.getInstance(project)
-            .getContentRootForFile(file.virtualFile) ?: return
+        globalContext.getUserData(ANALYZED_FILES)?.add(file)
+    }
 
-        val lazyResult = runReadAction {
-            BufAnalyzeUtils.checkLazily(project, disposable, contentRootForFile.toNioPath())
+    override fun inspectionFinished(
+        manager: InspectionManager,
+        globalContext: GlobalInspectionContext,
+        problemDescriptionsProcessor: ProblemDescriptionsProcessor,
+    ) {
+        val analyzedFiles = globalContext.getUserData(ANALYZED_FILES) ?: return
+        globalContext.putUserData(ANALYZED_FILES, null)
+        if (analyzedFiles.isEmpty()) return
+        val project = manager.project
+        val protoModificationTracker = PsiModificationTracker.getInstance(project).forLanguage(protobufLanguage())
+        repeat(MAX_ANALYZE_ATTEMPTS) { attempt ->
+            ProgressManager.checkCanceled()
+            val modificationCount = protoModificationTracker.modificationCount
+            val disposable = project.messageBus.createDisposableOnAnyPsiChange()
+                .also { Disposer.register(appService, it) }
+            val lazyResults = runReadAction {
+                val fileIndex = ProjectFileIndex.getInstance(project)
+                analyzedFiles.filter { it.isValid }
+                    .mapNotNull { fileIndex.getContentRootForFile(it.virtualFile) }
+                    .distinct()
+                    .map { BufAnalyzeUtils.checkLazily(project, disposable, it.toNioPath()) }
+            }
+            val futures = lazyResults.map { lazyResult ->
+                ApplicationManager.getApplication().executeOnPooledThread<BufAnalyzeResult?> { lazyResult.value }
+            }
+            val results = futures.mapNotNull { ProgressIndicatorUtils.awaitWithCheckCanceled(it) }
+            val reported = runReadAction {
+                // Proto PSI changed while buf was running, so offsets may no longer
+                // line up. Retry, but report on the last attempt rather than looping
+                // for as long as the user keeps editing.
+                val isLastAttempt = attempt == MAX_ANALYZE_ATTEMPTS - 1
+                if (!isLastAttempt && protoModificationTracker.modificationCount != modificationCount) {
+                    return@runReadAction false
+                }
+                for (file in analyzedFiles) {
+                    if (!file.isValid) continue
+                    for (result in results) {
+                        reportProblems(file, result, globalContext, problemDescriptionsProcessor)
+                    }
+                }
+                true
+            }
+            if (reported) return
         }
+    }
 
-        val result = ApplicationManager.getApplication().executeOnPooledThread<BufAnalyzeResult?> {
-            lazyResult.value
-        }.get() ?: return
-
+    private fun reportProblems(
+        file: PsiFile,
+        result: BufAnalyzeResult,
+        globalContext: GlobalInspectionContext,
+        problemDescriptionsProcessor: ProblemDescriptionsProcessor,
+    ) {
         val highlightInfos = mutableListOf<HighlightInfo>()
         highlightInfos.addHighlightsForFile(file, result)
 

src/main/resources/inspectionDescriptions/BufAnalyze.html

@@ -0,0 +1,5 @@
+<html>
+<body>
+Reports issues found by <code>buf lint</code> and <code>buf breaking</code>.
+</body>
+</html>

src/test/kotlin/build/buf/intellij/inspections/BufAnalyzeInspectionTest.kt

@@ -0,0 +1,61 @@
+// Copyright 2022-2026 Buf Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package build.buf.intellij.inspections
+
+import build.buf.intellij.base.BufTestBase
+import com.intellij.analysis.AnalysisScope
+import com.intellij.codeInspection.ex.GlobalInspectionToolWrapper
+import com.intellij.openapi.command.WriteCommandAction
+import com.intellij.openapi.fileEditor.FileDocumentManager
+import com.intellij.psi.PsiDocumentManager
+import com.intellij.testFramework.InspectionTestUtil
+import com.intellij.testFramework.createGlobalContextForTool
+
+class BufAnalyzeInspectionTest : BufTestBase() {
+    override fun getBasePath(): String = "inspections"
+
+    fun testUnsavedDocument() {
+        myFixture.configureByText(
+            "snake_case.proto",
+            """
+            syntax = "proto3";
+
+            message Foo {
+              string bar = 1;
+            }
+            """.trimIndent(),
+        )
+        // An unsaved edit forces the analysis to save documents (a write action)
+        // before running buf, which previously deadlocked batch inspections.
+        val document = myFixture.editor.document
+        WriteCommandAction.runWriteCommandAction(project) {
+            val offset = document.text.indexOf("bar")
+            document.replaceString(offset, offset + "bar".length, "Bar")
+        }
+        PsiDocumentManager.getInstance(project).commitAllDocuments()
+        assertTrue(FileDocumentManager.getInstance().isDocumentUnsaved(document))
+
+        val toolWrapper = GlobalInspectionToolWrapper(BufAnalyzeInspection())
+        val scope = AnalysisScope(myFixture.file)
+        val globalContext = createGlobalContextForTool(scope, project, listOf(toolWrapper))
+        InspectionTestUtil.runTool(toolWrapper, scope, globalContext)
+        InspectionTestUtil.compareToolResults(
+            globalContext,
+            toolWrapper,
+            false,
+            findTestDataFolder().resolve("unsavedDocument").toString(),
+        )
+    }
+}

src/test/resources/testData/inspections/unsavedDocument/expected.xml

@@ -0,0 +1,7 @@
+<problems>
+  <problem>
+    <file>snake_case.proto</file>
+    <line>4</line>
+    <description>Field name "Bar" should be lower_snake_case, such as "bar".</description>
+  </problem>
+</problems>