diff --git a/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt b/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt index 8c0be0a0599c..d4603f05aef1 100644 --- a/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt +++ b/platform/platform-impl/codeinsight-inline/src/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainer.kt @@ -72,8 +72,11 @@ class InlineCompletionLogsContainer() { private val asyncAdds = ConcurrentLinkedQueue() private suspend fun awaitAllAlreadyRunningAsyncAdds() { - while (currentCoroutineContext().isActive) { - val job = asyncAdds.poll() ?: return + // A snapshot on purpose: wait only for the adds that are already running, and never remove a job here. + // If this coroutine is canceled inside `join`, the job must stay visible to [cancelAsyncAdds], otherwise it keeps running + // on the application-level scope and retains everything its block captured (LLM-17026). + for (job in asyncAdds.toList()) { + if (!currentCoroutineContext().isActive) return job.join() } } @@ -107,6 +110,9 @@ class InlineCompletionLogsContainer() { block().forEach { add(it) } } asyncAdds.add(job) + // [asyncAdds] is a registry of *pending* jobs: a finished job is useless for [cancelAsyncAdds], and only the producer may + // remove one, never an awaiter (see [awaitAllAlreadyRunningAsyncAdds]). + job.invokeOnCompletion { asyncAdds.remove(job) } } /** diff --git a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt index b5808f661424..715b4d57545b 100644 --- a/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt +++ b/platform/platform-tests/testSrc/com/intellij/codeInsight/inline/completion/logs/InlineCompletionLogsContainerTest.kt @@ -10,6 +10,11 @@ import com.intellij.testFramework.ExtensionTestUtil import com.intellij.testFramework.LeakHunter import com.intellij.testFramework.LightPlatformTestCase import com.intellij.testFramework.common.timeoutRunBlocking +import kotlinx.coroutines.CompletableDeferred +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.awaitCancellation +import kotlinx.coroutines.cancelAndJoin +import kotlinx.coroutines.launch import org.junit.Test import org.junit.runner.RunWith import org.junit.runners.JUnit4 @@ -194,6 +199,40 @@ class InlineCompletionLogsContainerTest : LightPlatformTestCase() { LeakHunter.checkLeak(logsContainer, Project::class.java) } + /** + * A canceled awaiter must not drop a still running async add from the cancellation registry: + * [InlineCompletionLogsContainer.logCurrent] must still be able to cancel it, otherwise the add keeps running on the + * application-level scope and retains everything its block captured. See LLM-17026. + */ + @Test + fun testCanceledAwaitDoesNotStrandAsyncAdd(): Unit = timeoutRunBlocking { + val logsContainer = InlineCompletionLogsContainer() + logsContainer.mockRandom(1f) + val started = CompletableDeferred() + val blockCanceled = CompletableDeferred() + logsContainer.addAsync { + started.complete(Unit) + try { + awaitCancellation() + } + finally { + blockCanceled.complete(Unit) + } + } + started.await() + + try { + // `Dispatchers.Unconfined` runs the awaiter eagerly on this thread up to its first real suspension, + // so it is guaranteed to be inside `join` by the time `launch` returns. + val awaiter = launch(Dispatchers.Unconfined) { logsContainer.awaitAndGetCurrentLogs() } + awaiter.cancelAndJoin() + } + finally { + logsContainer.logCurrent(project = null) + } + blockCanceled.await() + } + private fun withEap(isEAP: Boolean, action: () -> Unit) { try { InlineCompletionEapSupport.getInstance().setMockEap(isEAP)