[Gradle] KotlinPluginLifecycle: requiredStage: Wait for stage if necessary

^KT-58676 In Progress

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)
 }
 

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 {