From b9a8f6dce2622424d95a85b68f37b02d752c6a4d Mon Sep 17 00:00:00 2001
From: Sebastian Sellmair <sebastian.sellmair@jetbrains.com>
Date: Wed, 21 Jun 2023 10:56:07 +0200
Subject: [PATCH] [Gradle] KotlinPluginLifecycle: requiredStage: Wait for stage
 if necessary

^KT-58676 In Progress
---
 .../kotlin/gradle/plugin/KotlinPluginLifecycle.kt      |  2 ++
 .../gradle/unitTests/KotlinPluginLifecycleTest.kt      | 10 ++++++++++
 2 files changed, 12 insertions(+)

diff --git a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt
index 548203e4776..3acdf1a8850 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt
+++ b/libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/plugin/KotlinPluginLifecycle.kt
@@ -191,8 +191,10 @@ internal suspend fun <T : Any> Property<T>.awaitFinalValue(): T? {
  * See also [withRestrictedStages]
  *
  * Will ensure that the given [block] can only execute in the given [stage]
+ * Will wait for the given [stage] if not arrived yet
  */
 internal suspend fun <T> requiredStage(stage: Stage, block: suspend () -> T): T {
+    if (currentKotlinPluginLifecycle().stage < stage) stage.await()
     return withRestrictedStages(hashSetOf(stage), block)
 }
 
diff --git a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt
index b8227ca0d17..64e2cc87545 100644
--- a/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt
+++ b/libraries/tools/kotlin-gradle-plugin/src/functionalTest/kotlin/org/jetbrains/kotlin/gradle/unitTests/KotlinPluginLifecycleTest.kt
@@ -458,6 +458,16 @@ class KotlinPluginLifecycleTest {
         }
     }
 
+    @Test
+    fun `test - requireStage`() = project.runLifecycleAwareTest {
+        launch {
+            requiredStage(AfterFinaliseDsl) {
+                assertEquals(AfterFinaliseDsl, stage)
+                assertFailsWith<IllegalLifecycleException> { AfterFinaliseDsl.nextOrThrow.await() }
+            }
+        }
+    }
+
     @Test
     fun `test - withRestrictedStages`() = project.runLifecycleAwareTest {
         launch {