Dispose AnalysisEnvironment (#2755)

kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/AnalysisContext.kt

@@ -0,0 +1,94 @@
+package org.jetbrains.dokka.analysis
+
+import org.jetbrains.dokka.DokkaConfiguration
+import org.jetbrains.dokka.Platform
+import org.jetbrains.dokka.utilities.DokkaLogger
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSeverity
+import org.jetbrains.kotlin.cli.common.messages.CompilerMessageSourceLocation
+import org.jetbrains.kotlin.cli.common.messages.MessageCollector
+import org.jetbrains.kotlin.cli.common.messages.MessageRenderer
+import org.jetbrains.kotlin.cli.jvm.compiler.KotlinCoreEnvironment
+import java.io.Closeable
+import java.io.File
+
+internal fun createAnalysisContext(
+    logger: DokkaLogger,
+    sourceSets: List<DokkaConfiguration.DokkaSourceSet>,
+    sourceSet: DokkaConfiguration.DokkaSourceSet,
+    analysisConfiguration: DokkaAnalysisConfiguration
+): AnalysisContext {
+    val parentSourceSets = sourceSets.filter { it.sourceSetID in sourceSet.dependentSourceSets }
+    val classpath = sourceSet.classpath + parentSourceSets.flatMap { it.classpath }
+    val sources = sourceSet.sourceRoots + parentSourceSets.flatMap { it.sourceRoots }
+
+    return createAnalysisContext(
+        logger = logger,
+        classpath = classpath,
+        sourceRoots = sources,
+        sourceSet = sourceSet,
+        analysisConfiguration = analysisConfiguration
+    )
+}
+
+internal fun createAnalysisContext(
+    logger: DokkaLogger,
+    classpath: List<File>,
+    sourceRoots: Set<File>,
+    sourceSet: DokkaConfiguration.DokkaSourceSet,
+    analysisConfiguration: DokkaAnalysisConfiguration
+): AnalysisContext {
+    val analysisEnvironment = AnalysisEnvironment(DokkaMessageCollector(logger), sourceSet.analysisPlatform).apply {
+        if (analysisPlatform == Platform.jvm) {
+            configureJdkClasspathRoots()
+        }
+        addClasspath(classpath)
+        addSources(sourceRoots)
+
+        loadLanguageVersionSettings(sourceSet.languageVersion, sourceSet.apiVersion)
+    }
+
+    val environment = analysisEnvironment.createCoreEnvironment()
+    val (facade, _) = analysisEnvironment.createResolutionFacade(
+        environment,
+        analysisConfiguration.ignoreCommonBuiltIns
+    )
+
+    return AnalysisContext(environment, facade, analysisEnvironment)
+}
+
+class DokkaMessageCollector(private val logger: DokkaLogger) : MessageCollector {
+    override fun clear() {
+        seenErrors = false
+    }
+
+    private var seenErrors = false
+
+    override fun report(severity: CompilerMessageSeverity, message: String, location: CompilerMessageSourceLocation?) {
+        if (severity == CompilerMessageSeverity.ERROR) {
+            seenErrors = true
+        }
+        logger.info(MessageRenderer.PLAIN_FULL_PATHS.render(severity, message, location))
+    }
+
+    override fun hasErrors() = seenErrors
+}
+
+// It is not data class due to ill-defined equals
+class AnalysisContext(
+    environment: KotlinCoreEnvironment,
+    facade: DokkaResolutionFacade,
+    private val analysisEnvironment: AnalysisEnvironment
+) : Closeable {
+    private var isClosed: Boolean = false
+    val environment: KotlinCoreEnvironment = environment
+        get() = field.takeUnless { isClosed } ?: throw IllegalStateException("AnalysisEnvironment is already closed")
+    val facade: DokkaResolutionFacade = facade
+        get() = field.takeUnless { isClosed } ?: throw IllegalStateException("AnalysisEnvironment is already closed")
+
+    operator fun component1() = environment
+    operator fun component2() = facade
+    override fun close() {
+        isClosed = true
+        analysisEnvironment.dispose()
+    }
+}

kotlin-analysis/src/main/kotlin/org/jetbrains/dokka/analysis/KotlinAnalysis.kt

@@ -7,18 +7,48 @@ import org.jetbrains.dokka.DokkaSourceSetID
 import org.jetbrains.dokka.model.SourceSetDependent
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.utilities.DokkaLogger
+import java.io.Closeable
 
-fun KotlinAnalysis(sourceSets: List<DokkaSourceSet>, logger: DokkaLogger, analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()): KotlinAnalysis {
+fun ProjectKotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+): KotlinAnalysis {
     val environments = sourceSets.associateWith { sourceSet ->
-        createEnvironmentAndFacade(
+        createAnalysisContext(
             logger = logger,
             sourceSets = sourceSets,
             sourceSet = sourceSet,
             analysisConfiguration = analysisConfiguration
         )
     }
+    return EnvironmentKotlinAnalysis(environments)
+}
+
+/**
+ * [projectKotlinAnalysis] needs to be closed separately
+ *  Usually the analysis created for samples is short-lived and can be closed right after
+ *  it's been used, there's no need to wait for [projectKotlinAnalysis] to be closed as it must be handled separately.
+ */
+fun SamplesKotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    projectKotlinAnalysis: KotlinAnalysis,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+): KotlinAnalysis {
+    val environments = sourceSets
+        .filter { it.samples.isNotEmpty() }
+        .associateWith { sourceSet ->
+            createAnalysisContext(
+                logger = logger,
+                classpath = sourceSet.classpath,
+                sourceRoots = sourceSet.samples,
+                sourceSet = sourceSet,
+                analysisConfiguration = analysisConfiguration
+            )
+    }
 
-    return KotlinAnalysisImpl(environments)
+    return EnvironmentKotlinAnalysis(environments, projectKotlinAnalysis)
 }
 
 class DokkaAnalysisConfiguration(
@@ -32,22 +62,48 @@ class DokkaAnalysisConfiguration(
 @Deprecated(message = "Construct using list of DokkaSourceSets and logger",
     replaceWith = ReplaceWith("KotlinAnalysis(context.configuration.sourceSets, context.logger)")
 )
-fun KotlinAnalysis(context: DokkaContext): KotlinAnalysis = KotlinAnalysis(context.configuration.sourceSets, context.logger)
+fun KotlinAnalysis(context: DokkaContext): KotlinAnalysis =
+    ProjectKotlinAnalysis(context.configuration.sourceSets, context.logger)
+
+@Deprecated(message = "It was renamed to `ProjectKotlinAnalysis`",
+    replaceWith = ReplaceWith("ProjectKotlinAnalysis(sourceSets, logger, analysisConfiguration)")
+)
+fun KotlinAnalysis(
+    sourceSets: List<DokkaSourceSet>,
+    logger: DokkaLogger,
+    analysisConfiguration: DokkaAnalysisConfiguration = DokkaAnalysisConfiguration()
+) = ProjectKotlinAnalysis(sourceSets, logger, analysisConfiguration)
 
-interface KotlinAnalysis : SourceSetDependent<EnvironmentAndFacade> {
-    override fun get(key: DokkaSourceSet): EnvironmentAndFacade
-    operator fun get(sourceSetID: DokkaSourceSetID): EnvironmentAndFacade
-}
 
-internal class KotlinAnalysisImpl(
-    private val environments: SourceSetDependent<EnvironmentAndFacade>
-) : KotlinAnalysis, SourceSetDependent<EnvironmentAndFacade> by environments {
+/**
+ * First child delegation. It does not close [parent].
+ */
+abstract class KotlinAnalysis(
+    val parent: KotlinAnalysis? = null
+) : Closeable {
 
-    override fun get(key: DokkaSourceSet): EnvironmentAndFacade {
-        return environments[key] ?: throw IllegalStateException("Missing EnvironmentAndFacade for sourceSet $key")
+    operator fun get(key: DokkaSourceSet): AnalysisContext {
+        return get(key.sourceSetID)
     }
+    operator fun get(key: DokkaSourceSetID): AnalysisContext {
+        return find(key)
+            ?: parent?.get(key)
+            ?: throw IllegalStateException("Missing EnvironmentAndFacade for sourceSet ${key}")
+    }
+    protected abstract fun find(sourceSetID: DokkaSourceSetID): AnalysisContext?
+}
+
+internal open class EnvironmentKotlinAnalysis(
+    private val environments: SourceSetDependent<AnalysisContext>,
+    parent: KotlinAnalysis? = null,
+) : KotlinAnalysis(parent = parent) {
 
-    override fun get(sourceSetID: DokkaSourceSetID): EnvironmentAndFacade {
-        return environments.entries.first { (sourceSet, _) -> sourceSet.sourceSetID == sourceSetID }.value
+    override fun find(sourceSetID: DokkaSourceSetID): AnalysisContext? =
+        environments.entries.firstOrNull { (sourceSet, _) -> sourceSet.sourceSetID == sourceSetID }?.value
+
+    override fun close() {
+        environments.values.forEach(AnalysisContext::close)
     }
 }
+
+

plugins/base/base-test-utils/src/main/kotlin/testRunner/baseTestApi.kt

@@ -53,6 +53,8 @@ class BaseDokkaTestGenerator(
         singleModuleGeneration.render(transformedPages)
         renderingStage(transformedPages, context)
 
+        singleModuleGeneration.runPostActions()
+
         singleModuleGeneration.reportAfterRendering()
     }
 }

plugins/base/src/main/kotlin/DokkaBase.kt

@@ -4,6 +4,7 @@ package org.jetbrains.dokka.base
 
 import org.jetbrains.dokka.CoreExtensions
 import org.jetbrains.dokka.analysis.KotlinAnalysis
+import org.jetbrains.dokka.analysis.ProjectKotlinAnalysis
 import org.jetbrains.dokka.base.renderers.*
 import org.jetbrains.dokka.base.renderers.html.*
 import org.jetbrains.dokka.base.renderers.html.command.consumers.PathToRootConsumer
@@ -36,6 +37,8 @@ import org.jetbrains.dokka.base.translators.descriptors.ExternalClasslikesTransl
 import org.jetbrains.dokka.base.translators.descriptors.ExternalDocumentablesProvider
 import org.jetbrains.dokka.base.utils.NoopIntellijLoggerFactory
 import org.jetbrains.dokka.plugability.DokkaPlugin
+import org.jetbrains.dokka.plugability.querySingle
+import org.jetbrains.dokka.renderers.PostAction
 import org.jetbrains.dokka.transformers.documentation.PreMergeDocumentableTransformer
 import org.jetbrains.dokka.transformers.pages.PageTransformer
 
@@ -189,7 +192,7 @@ class DokkaBase : DokkaPlugin() {
 
     val defaultKotlinAnalysis by extending {
         kotlinAnalysis providing { ctx ->
-            KotlinAnalysis(
+            ProjectKotlinAnalysis(
                 sourceSets = ctx.configuration.sourceSets,
                 logger = ctx.logger
             )
@@ -281,6 +284,10 @@ class DokkaBase : DokkaPlugin() {
         externalClasslikesTranslator providing ::DefaultDescriptorToDocumentableTranslator
     }
 
+    internal val disposeKotlinAnalysisPostAction by extending {
+        CoreExtensions.postActions with PostAction { this@DokkaBase.querySingle { kotlinAnalysis }.close() }
+    }
+
     private companion object {
         init {
             // Suppress messages emitted by the IntelliJ logger since

plugins/base/src/main/kotlin/transformers/pages/samples/SamplesTransformer.kt

@@ -4,11 +4,7 @@ import com.intellij.psi.PsiElement
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
 import org.jetbrains.dokka.DokkaConfiguration.DokkaSourceSet
-import org.jetbrains.dokka.Platform
-import org.jetbrains.dokka.analysis.AnalysisEnvironment
-import org.jetbrains.dokka.analysis.DokkaMessageCollector
-import org.jetbrains.dokka.analysis.DokkaResolutionFacade
-import org.jetbrains.dokka.analysis.EnvironmentAndFacade
+import org.jetbrains.dokka.analysis.*
 import org.jetbrains.dokka.base.DokkaBase
 import org.jetbrains.dokka.base.renderers.sourceSets
 import org.jetbrains.dokka.links.DRI
@@ -38,51 +34,36 @@ abstract class SamplesTransformer(val context: DokkaContext) : PageTransformer {
          * Currently, all `ThreadLocal`s are in a compiler/IDE codebase.
          */
         runBlocking(Dispatchers.Default) {
-            val analysis = setUpAnalysis(context)
-
-            input.transformContentPagesTree { page ->
-                val samples = (page as? WithDocumentables)?.documentables?.flatMap {
-                    it.documentation.entries.flatMap { entry ->
-                        entry.value.children.filterIsInstance<Sample>().map { entry.key to it }
+            val analysis = SamplesKotlinAnalysis(
+                sourceSets = context.configuration.sourceSets,
+                logger = context.logger,
+                projectKotlinAnalysis = context.plugin<DokkaBase>().querySingle { kotlinAnalysis }
+            )
+            analysis.use {
+                input.transformContentPagesTree { page ->
+                    val samples = (page as? WithDocumentables)?.documentables?.flatMap {
+                        it.documentation.entries.flatMap { entry ->
+                            entry.value.children.filterIsInstance<Sample>().map { entry.key to it }
+                        }
                     }
-                }
-
-                samples?.fold(page as ContentPage) { acc, (sampleSourceSet, sample) ->
-                    acc.modified(
-                        content = acc.content.addSample(page, sampleSourceSet, sample.name, analysis),
-                        embeddedResources = acc.embeddedResources + KOTLIN_PLAYGROUND_SCRIPT
-                    )
-                } ?: page
-            }
-        }
 
-    private fun setUpAnalysis(context: DokkaContext) = context.configuration.sourceSets.associateWith { sourceSet ->
-        if (sourceSet.samples.isEmpty()) context.plugin<DokkaBase>()
-            .querySingle { kotlinAnalysis }[sourceSet] // from sourceSet.sourceRoots
-        else AnalysisEnvironment(DokkaMessageCollector(context.logger), sourceSet.analysisPlatform).run {
-            if (analysisPlatform == Platform.jvm) {
-                configureJdkClasspathRoots()
+                    samples?.fold(page as ContentPage) { acc, (sampleSourceSet, sample) ->
+                        acc.modified(
+                            content = acc.content.addSample(page, sampleSourceSet, sample.name, it),
+                            embeddedResources = acc.embeddedResources + KOTLIN_PLAYGROUND_SCRIPT
+                        )
+                    } ?: page
+                }
             }
-            sourceSet.classpath.forEach(::addClasspath)
-
-            addSources(sourceSet.samples.toList())
-
-            loadLanguageVersionSettings(sourceSet.languageVersion, sourceSet.apiVersion)
-
-            val environment = createCoreEnvironment()
-            val (facade, _) = createResolutionFacade(environment)
-            EnvironmentAndFacade(environment, facade)
         }
-    }
 
     private fun ContentNode.addSample(
         contentPage: ContentPage,
         sourceSet: DokkaSourceSet,
         fqName: String,
-        analysis: Map<DokkaSourceSet, EnvironmentAndFacade>
+        analysis: KotlinAnalysis
     ): ContentNode {
-        val facade = analysis[sourceSet]?.facade
-            ?: return this.also { context.logger.warn("Cannot resolve facade for platform ${sourceSet.sourceSetID}") }
+        val facade = analysis[sourceSet].facade
         val psiElement = fqNameToPsiElement(facade, fqName)
             ?: return this.also { context.logger.warn("Cannot find PsiElement corresponding to $fqName") }
         val imports =