[Gradle] invokeWhenCreated: Lenient until last stage of lifecycle

^KT-34662 Verification Pending

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

@@ -32,6 +32,7 @@ import org.jetbrains.kotlin.konan.target.CompilerOutputKind
 import org.jetbrains.kotlin.statistics.metrics.StringMetrics
 import java.util.concurrent.atomic.AtomicBoolean
 import javax.inject.Inject
+import javax.inject.Provider
 import kotlin.reflect.KClass
 
 private const val KOTLIN_PROJECT_EXTENSION_NAME = "kotlin"
@@ -134,14 +135,10 @@ abstract class KotlinTopLevelExtension(internal val project: Project) : KotlinTo
      */
     @ExperimentalKotlinGradlePluginApi
     fun <T : Named> NamedDomainObjectContainer<T>.invokeWhenCreated(name: String, configure: T.() -> Unit) {
-        val invoked = AtomicBoolean(false)
-        matching { it.name == name }.all { value -> if (!invoked.getAndSet(true)) value.configure() }
-        project.whenEvaluated {
-            if (!invoked.getAndSet(true)) {
-                /*
-                Expected to fail, since the listener was not called.
-                Letting Gradle fail here will result in a more natural error message
-                */
+        configureEach { if (it.name == name) it.configure() }
+        project.launchInStage(KotlinPluginLifecycle.Stage.ReadyForExecution) {
+            if (name !in names) {
+                /* Expect 'named' to throw corresponding exception */
                 named(name).configure(configure)
             }
         }

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

@@ -10,9 +10,13 @@ package org.jetbrains.kotlin.gradle.unitTests
 import org.gradle.api.Named
 import org.gradle.api.UnknownDomainObjectException
 import org.jetbrains.kotlin.gradle.dsl.kotlinExtension
+import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle
+import org.jetbrains.kotlin.gradle.plugin.await
 import org.jetbrains.kotlin.gradle.util.buildProjectWithJvm
+import org.jetbrains.kotlin.gradle.util.runLifecycleAwareTest
 import org.jetbrains.kotlin.tooling.core.withLinearClosure
 import org.junit.Test
+import java.util.concurrent.atomic.AtomicInteger
 import kotlin.test.assertEquals
 import kotlin.test.assertFails
 import kotlin.test.fail
@@ -68,4 +72,16 @@ class InvokeWhenCreatedTest {
         val causes = assertFails { project.evaluate() }.withLinearClosure { it.cause }
         assertEquals<Class<*>>(UnknownDomainObjectException::class.java, causes.last().javaClass)
     }
+
+    @Test
+    fun `test - is lenient up-to last stage of Kotlin Plugin Lifecycle`() = project.runLifecycleAwareTest {
+        project.kotlinExtension.apply {
+            val invocations = AtomicInteger(0)
+            container.invokeWhenCreated("x") { assertEquals(1, invocations.incrementAndGet()) }
+            await(KotlinPluginLifecycle.Stage.last)
+            assertEquals(0, invocations.get())
+            container.create("x")
+            assertEquals(1, invocations.get())
+        }
+    }
 }