[Gradle] KotlinMultiplatformPluginLifecycle: Support suspend and await

KT-34662

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinMultiplatformPluginLifecycle.kt

@@ -10,6 +10,7 @@ import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Ill
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Stage
 import org.jetbrains.kotlin.gradle.utils.getOrPut
 import java.util.concurrent.atomic.AtomicBoolean
+import kotlin.coroutines.*
 
 internal val Project.kotlinMultiplatformPluginLifecycle: KotlinMultiplatformPluginLifecycle
     get() = extraProperties.getOrPut(KotlinMultiplatformPluginLifecycle::class.java.name) { KotlinMultiplatformPluginLifecycleImpl() }
@@ -26,8 +27,15 @@ internal interface KotlinMultiplatformPluginLifecycle {
     enum class Stage {
         Configure,
         AfterEvaluate,
+
+        BeforeFinaliseRefinesEdges,
         FinaliseRefinesEdges,
+        AfterFinaliseRefinesEdges,
+
+        BeforeFinaliseCompilations,
         FinaliseCompilations,
+        AfterFinaliseCompilations,
+
         Finalised
     }
 
@@ -35,14 +43,19 @@ internal interface KotlinMultiplatformPluginLifecycle {
 
     fun enqueue(stage: Stage, action: () -> Unit)
 
+    fun launch(block: suspend KotlinMultiplatformPluginLifecycle.() -> Unit)
+
+    suspend fun await(stage: Stage)
+
     class IllegalLifecycleException(message: String) : IllegalStateException(message)
 }
 
 private class KotlinMultiplatformPluginLifecycleImpl : KotlinMultiplatformPluginLifecycle {
-    private val enqueuedStages = ArrayDeque(Stage.values().toList())
-    private val enqueuedActions = Stage.values().associateWith { ArrayDeque<() -> Unit>() }
-    private var configureLoopRunning = false
+    private val enqueuedStages: ArrayDeque<Stage> = ArrayDeque(Stage.values().toList())
+    private val enqueuedActions: Map<Stage, ArrayDeque<() -> Unit>> = Stage.values().associateWith { ArrayDeque() }
+    private var configureLoopRunning = AtomicBoolean(false)
     private var isStarted = AtomicBoolean(false)
+    private var isFinished = AtomicBoolean(false)
 
     fun start(project: Project) {
         check(!isStarted.getAndSet(true)) {
@@ -65,7 +78,10 @@ private class KotlinMultiplatformPluginLifecycleImpl : KotlinMultiplatformPlugin
             val action = queue.removeFirstOrNull()
             action?.invoke()
         } while (action != null)
-        val nextStage = enqueuedStages.removeFirstOrNull() ?: return
+        val nextStage = enqueuedStages.removeFirstOrNull() ?: run {
+            isFinished.set(true)
+            return
+        }
 
         project.afterEvaluate {
             executeStage(project, nextStage)
@@ -74,8 +90,7 @@ private class KotlinMultiplatformPluginLifecycleImpl : KotlinMultiplatformPlugin
 
     private fun startConfigureLoopIfNecessary() {
         check(stage == Stage.Configure) { "Cannot start 'configure loop' on stage '$stage'" }
-        if (configureLoopRunning) return
-        configureLoopRunning = true
+        if (configureLoopRunning.getAndSet(true)) return
         try {
             val queue = enqueuedActions.getValue(Stage.Configure)
             do {
@@ -83,7 +98,7 @@ private class KotlinMultiplatformPluginLifecycleImpl : KotlinMultiplatformPlugin
                 action?.invoke()
             } while (action != null)
         } finally {
-            configureLoopRunning = false
+            configureLoopRunning.set(false)
         }
     }
 
@@ -100,4 +115,26 @@ private class KotlinMultiplatformPluginLifecycleImpl : KotlinMultiplatformPlugin
             startConfigureLoopIfNecessary()
         }
     }
+
+    override fun launch(block: suspend KotlinMultiplatformPluginLifecycle.() -> Unit) {
+        check(isStarted.get()) { "Cannot launch when ${KotlinMultiplatformPluginLifecycle::class.simpleName} is not started" }
+        check(!isFinished.get()) { "Cannot launch when ${KotlinMultiplatformPluginLifecycle::class.simpleName} is already finished" }
+        val coroutine = block.createCoroutine(this, object : Continuation<Unit> {
+            override val context: CoroutineContext = EmptyCoroutineContext
+            override fun resumeWith(result: Result<Unit>) = result.getOrThrow()
+        })
+
+        enqueue(stage) {
+            coroutine.resume(Unit)
+        }
+    }
+
+    override suspend fun await(stage: Stage) {
+        if (this.stage >= stage) return
+        suspendCoroutine<Unit> { continuation ->
+            enqueue(stage) {
+                continuation.resume(Unit)
+            }
+        }
+    }
 }

libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinMultiplatformPluginLifecycleTest.kt

@@ -7,10 +7,13 @@
 
 package org.jetbrains.kotlin.gradle.unitTests
 
+import org.gradle.api.ProjectConfigurationException
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle
 import org.jetbrains.kotlin.gradle.plugin.KotlinMultiplatformPluginLifecycle.Stage.*
 import org.jetbrains.kotlin.gradle.plugin.kotlinMultiplatformPluginLifecycle
 import org.jetbrains.kotlin.gradle.util.buildProjectWithMPP
+import org.jetbrains.kotlin.tooling.core.linearClosure
+import org.jetbrains.kotlin.tooling.core.withLinearClosure
 import org.junit.Test
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
@@ -210,4 +213,95 @@ class KotlinMultiplatformPluginLifecycleTest {
         assertEquals(1, action2Invocations.get())
         assertEquals(1, action3Invocations.get())
     }
+
+    @Test
+    fun `test - launch in configure`() {
+        val action1Invocations = AtomicInteger(0)
+        val action2Invocations = AtomicInteger(0)
+        val action3Invocations = AtomicInteger(0)
+
+        lifecycle.launch action1@{
+            lifecycle.launch action2@{
+                assertEquals(1, action1Invocations.get())
+                assertEquals(0, action3Invocations.get())
+                assertEquals(1, action2Invocations.incrementAndGet())
+            }
+            assertEquals(0, action2Invocations.get())
+            assertEquals(0, action3Invocations.get())
+            assertEquals(1, action1Invocations.incrementAndGet())
+
+            lifecycle.launch action3@{
+                assertEquals(1, action1Invocations.get())
+                assertEquals(1, action2Invocations.get())
+                assertEquals(1, action3Invocations.incrementAndGet())
+            }
+        }
+
+        assertEquals(1, action1Invocations.get())
+        assertEquals(1, action2Invocations.get())
+        assertEquals(1, action3Invocations.get())
+    }
+
+    @Test
+    fun `test - launch in configure - await Stage`() {
+        val executionPointA = AtomicBoolean(false)
+        val executionPointB = AtomicBoolean(false)
+        lifecycle.launch action1@{
+            assertFalse(executionPointA.getAndSet(true))
+            assertEquals(Configure, stage)
+            await(AfterEvaluate)
+            assertEquals(AfterEvaluate, stage)
+            assertFalse(executionPointB.getAndSet(true))
+        }
+
+        assertTrue(executionPointA.get())
+        assertFalse(executionPointB.get())
+        project.evaluate()
+        assertTrue(executionPointA.get())
+        assertTrue(executionPointB.get())
+    }
+
+    @Test
+    fun `test - launch - await - launch`() {
+        val executedInnerAction = AtomicBoolean(false)
+        lifecycle.launch {
+            await(AfterEvaluate)
+            launch {
+                assertEquals(AfterEvaluate, stage)
+                await(FinaliseRefinesEdges)
+                assertEquals(FinaliseRefinesEdges, stage)
+                assertFalse(executedInnerAction.getAndSet(true))
+            }
+        }
+
+        assertFalse(executedInnerAction.get())
+        project.evaluate()
+        assertTrue(executedInnerAction.get())
+    }
+
+    @Test
+    fun `test - launch - exception`() {
+        assertFailsWith<IllegalStateException> {
+            lifecycle.launch {
+                throw IllegalStateException("42")
+            }
+        }
+    }
+
+    @Test
+    fun `test - launch - await - exception`() {
+        val testException = object : Throwable() {}
+        lifecycle.launch {
+            await(AfterEvaluate)
+            launch {
+                throw testException
+            }
+        }
+
+        val causes = assertFailsWith<ProjectConfigurationException> {
+            project.evaluate()
+        }.withLinearClosure<Throwable> { it.cause }
+
+        assertTrue(testException in causes)
+    }
 }