Fix memory leak caused by coroutines (#1751)

* Fix memory leak caused by coroutines

* Catch all exceptions that might fail the build

Co-authored-by: Marcin Aman <marcin.aman@gmail.com>

2 files changed, 41 insertions, 5 deletions

diff --git a/core/src/main/kotlin/DokkaGenerator.kt b/core/src/main/kotlin/DokkaGenerator.kt
index e3dcf626..9241d728 100644
--- a/core/src/main/kotlin/DokkaGenerator.kt
+++ b/core/src/main/kotlin/DokkaGenerator.kt
@@ -2,10 +2,15 @@
 
 package org.jetbrains.dokka
 
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.InternalCoroutinesApi
+import kotlinx.coroutines.scheduling.ExperimentalCoroutineDispatcher
 import org.jetbrains.dokka.generation.GracefulGenerationExit
 import org.jetbrains.dokka.plugability.DokkaContext
 import org.jetbrains.dokka.plugability.DokkaPlugin
 import org.jetbrains.dokka.utilities.DokkaLogger
+import kotlin.reflect.full.memberProperties
+import kotlin.reflect.jvm.isAccessible
 
 /**
  * DokkaGenerator is the main entry point for generating documentation
@@ -16,15 +21,22 @@ class DokkaGenerator(
     private val configuration: DokkaConfiguration,
     private val logger: DokkaLogger
 ) {
+
     fun generate() = timed(logger) {
         report("Initializing plugins")
         val context = initializePlugins(configuration, logger)
 
-        context.single(CoreExtensions.generation).run {
-            logger.progress("Dokka is performing: $generationName")
-            generate()
+        runCatching {
+            context.single(CoreExtensions.generation).run {
+                logger.progress("Dokka is performing: $generationName")
+                generate()
+            }
+        }.exceptionOrNull()?.let { e ->
+            finalizeCoroutines()
+            throw e
         }
 
+        finalizeCoroutines()
     }.dump("\n\n === TIME MEASUREMENT ===\n")
 
     fun initializePlugins(
@@ -32,6 +44,29 @@ class DokkaGenerator(
         logger: DokkaLogger,
         additionalPlugins: List<DokkaPlugin> = emptyList()
     ) = DokkaContext.create(configuration, logger, additionalPlugins)
+
+    private fun finalizeCoroutines() {
+        runCatching {
+            Dispatchers.Default.closeExecutor()
+
+            Dispatchers.IO.let { dispatcher ->
+                dispatcher::class.memberProperties.find {
+                    it.name == "dispatcher"
+                }?.also {
+                    it.isAccessible = true
+                }?.call(dispatcher)
+            }?.closeExecutor()
+        }
+    }
+
+    @OptIn(InternalCoroutinesApi::class)
+    private fun Any.closeExecutor() = (this as ExperimentalCoroutineDispatcher).also {
+        it.executor::class.members
+            .find { it.name == "close" }
+            ?.also {
+                it.isAccessible = true
+            }?.call(it.executor)
+    }
 }
 
 class Timer internal constructor(startTime: Long, private val logger: DokkaLogger?) {
diff --git a/core/src/main/kotlin/transformers/sources/AsyncSourceToDocumentableTranslator.kt b/core/src/main/kotlin/transformers/sources/AsyncSourceToDocumentableTranslator.kt
index 19113446..3c00be5a 100644
--- a/core/src/main/kotlin/transformers/sources/AsyncSourceToDocumentableTranslator.kt
+++ b/core/src/main/kotlin/transformers/sources/AsyncSourceToDocumentableTranslator.kt
@@ -1,5 +1,6 @@
 package org.jetbrains.dokka.transformers.sources
 
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.runBlocking
 import org.jetbrains.dokka.DokkaConfiguration
 import org.jetbrains.dokka.model.DModule
@@ -9,7 +10,7 @@ interface AsyncSourceToDocumentableTranslator : SourceToDocumentableTranslator {
     suspend fun invokeSuspending(sourceSet: DokkaConfiguration.DokkaSourceSet, context: DokkaContext): DModule
 
     override fun invoke(sourceSet: DokkaConfiguration.DokkaSourceSet, context: DokkaContext): DModule =
-        runBlocking {
+        runBlocking(Dispatchers.Default) {
             invokeSuspending(sourceSet, context)
         }
-}
\ No newline at end of file
+}