[Gradle] Split KotlinPluginLifecycle and KotlinPluginLifecycleImpl

^KT-58255 Verification Pending

Author: sellmair

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

@@ -8,13 +8,8 @@ package org.jetbrains.kotlin.gradle.plugin
 import org.gradle.api.Project
 import org.gradle.api.provider.Property
 import org.jetbrains.kotlin.gradle.plugin.KotlinPluginLifecycle.*
-import org.jetbrains.kotlin.gradle.utils.CompletableFuture
 import org.jetbrains.kotlin.gradle.utils.Future
-import org.jetbrains.kotlin.gradle.utils.failures
 import org.jetbrains.kotlin.gradle.utils.getOrPut
-import java.util.*
-import java.util.concurrent.atomic.AtomicBoolean
-import kotlin.collections.ArrayDeque
 import kotlin.coroutines.*
 
 /*
@@ -156,11 +151,7 @@ internal val Project.kotlinPluginLifecycle: KotlinPluginLifecycle
  * some later data to be available. In this case, the Future still will only return 'sane' data.
  */
 internal val Project.configurationResult: Future<ProjectConfigurationResult>
-    get() = configurationResultImpl
-
-
-private val Project.configurationResultImpl: CompletableFuture<ProjectConfigurationResult>
-    get() = extraProperties.getOrPut("org.jetbrains.kotlin.gradle.plugin.configurationResult") { CompletableFuture() }
+    get() = (kotlinPluginLifecycle as KotlinPluginLifecycleImpl).configurationResult
 
 
 /**
@@ -175,8 +166,7 @@ internal fun Project.startKotlinPluginLifecycle() {
  * the currently running coroutine. Throws if this coroutine was not started using a [KotlinPluginLifecycle]
  */
 internal suspend fun currentKotlinPluginLifecycle(): KotlinPluginLifecycle {
-    return coroutineContext[KotlinPluginLifecycleCoroutineContextElement]?.lifecycle
-        ?: error("Missing $KotlinPluginLifecycleCoroutineContextElement in currentCoroutineContext")
+    return coroutineContext.kotlinPluginLifecycle
 }
 
 /**
@@ -224,33 +214,6 @@ internal suspend fun <T> requireCurrentStage(block: suspend () -> T): T {
     return requiredStage(currentKotlinPluginLifecycle().stage, block)
 }
 
-/**
- * Will ensure that the given [block] cannot leave the specified allowed stages [allowed]
- * e.g.
- *
- * ```kotlin
- * project.launchInStage(Stage.BeforeFinaliseDsl) {
- *     withRestrictedStages(Stage.upTo(Stage.FinaliseDsl)) {
- *        await(Stage.FinaliseDsl) // <- OK, since still in allowed stages
- *        await(Stage.AfterFinaliseDsl) // <- fails, since not in allowed stages!
- *     }
- * }
- * ```
- */
-internal suspend fun <T> withRestrictedStages(allowed: Set<Stage>, block: suspend () -> T): T {
-    val newCoroutineContext = coroutineContext + RestrictedLifecycleStages(currentKotlinPluginLifecycle(), allowed)
-    return suspendCoroutine { continuation ->
-        val newContinuation = object : Continuation<T> {
-            override val context: CoroutineContext
-                get() = newCoroutineContext
-
-            override fun resumeWith(result: Result<T>) {
-                continuation.resumeWith(result)
-            }
-        }
-        block.startCoroutine(newContinuation)
-    }
-}
 
 /*
 Definition of the Lifecycle and its stages
@@ -339,195 +302,3 @@ internal interface KotlinPluginLifecycle {
 
     class IllegalLifecycleException(message: String) : IllegalStateException(message)
 }
-
-
-/*
-Implementation
- */
-
-internal class KotlinPluginLifecycleImpl(override val project: Project) : KotlinPluginLifecycle {
-    private val enqueuedActions: Map<Stage, ArrayDeque<KotlinPluginLifecycle.() -> Unit>> =
-        Stage.values().associateWith { ArrayDeque() }
-
-    private val loopRunning = AtomicBoolean(false)
-    private val isStarted = AtomicBoolean(false)
-    private val isFinishedSuccessfully = AtomicBoolean(false)
-    private val isFinishedWithFailures = AtomicBoolean(false)
-
-    override var stage: Stage = Stage.values.first()
-
-    fun start() {
-        check(!isStarted.getAndSet(true)) {
-            "${KotlinPluginLifecycle::class.java.name} already started"
-        }
-
-        check(!project.state.executed) {
-            "${KotlinPluginLifecycle::class.java.name} cannot be started in ProjectState '${project.state}'"
-        }
-
-        loopIfNecessary()
-
-        project.whenEvaluated {
-            /* Check for failures happening during buildscript evaluation */
-            project.failures.let { failures ->
-                if (failures.isNotEmpty()) {
-                    finishWithFailures(failures)
-                    return@whenEvaluated
-                }
-            }
-
-            assert(enqueuedActions.getValue(stage).isEmpty()) { "Expected empty queue from '$stage'" }
-            stage = stage.nextOrThrow
-            executeCurrentStageAndScheduleNext()
-        }
-    }
-
-    private fun executeCurrentStageAndScheduleNext() {
-        stage.previousOrNull?.let { previousStage ->
-            assert(enqueuedActions.getValue(previousStage).isEmpty()) {
-                "Actions from previous stage '$previousStage' have not been executed (stage: '$stage')"
-            }
-        }
-
-        val failures = project.failures
-        if (failures.isNotEmpty()) {
-            finishWithFailures(failures)
-            return
-        }
-
-        try {
-            loopIfNecessary()
-        } catch (t: Throwable) {
-            finishWithFailures(listOf(t))
-            throw t
-        }
-
-        stage = stage.nextOrNull ?: run {
-            finishSuccessfully()
-            return
-        }
-
-        project.afterEvaluate {
-            executeCurrentStageAndScheduleNext()
-        }
-    }
-
-    private fun loopIfNecessary() {
-        if (loopRunning.getAndSet(true)) return
-        try {
-            val queue = enqueuedActions.getValue(stage)
-            do {
-                project.state.rethrowFailure()
-                val action = queue.removeFirstOrNull()
-                action?.invoke(this)
-            } while (action != null)
-        } finally {
-            loopRunning.set(false)
-        }
-    }
-
-    private fun finishWithFailures(failures: List<Throwable>) {
-        assert(failures.isNotEmpty())
-        assert(isStarted.get())
-        assert(!isFinishedWithFailures.getAndSet(true))
-        project.configurationResultImpl.complete(ProjectConfigurationResult.Failure(failures))
-    }
-
-    private fun finishSuccessfully() {
-        assert(isStarted.get())
-        assert(!isFinishedSuccessfully.getAndSet(true))
-        project.configurationResultImpl.complete(ProjectConfigurationResult.Success)
-    }
-
-    fun enqueue(stage: Stage, action: KotlinPluginLifecycle.() -> Unit) {
-        if (stage < this.stage) {
-            throw IllegalLifecycleException("Cannot enqueue Action for stage '$stage' in current stage '${this.stage}'")
-        }
-
-        /*
-        Lifecycle finished: action shall not be enqueued, but just executed right away.
-        This is desirable, so that .enqueue (and .launch) functions that are scheduled in execution phase
-        will be executed right away (no suspend necessary or wanted)
-        */
-        if (isFinishedSuccessfully.get()) {
-            return action()
-        }
-
-        /*
-        Lifecycle finished, but some exceptions have been thrown.
-        In this case, an enqueue for future Stages is not allowed, since those will not be executed anymore.
-        Any enqueue in the current stage will be executed right away (no suspend necessary or wanted).
-         */
-        if (isFinishedWithFailures.get()) {
-            return if (stage == this.stage) action()
-            else Unit
-        }
-
-        enqueuedActions.getValue(stage).addLast(action)
-
-        if (stage == Stage.EvaluateBuildscript && isStarted.get()) {
-            loopIfNecessary()
-        }
-    }
-
-    override fun launch(block: suspend KotlinPluginLifecycle.() -> Unit) {
-        val lifecycle = this
-
-        val coroutine = block.createCoroutine(this, object : Continuation<Unit> {
-            override val context: CoroutineContext = EmptyCoroutineContext +
-                    KotlinPluginLifecycleCoroutineContextElement(lifecycle)
-
-            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)
-            }
-        }
-    }
-}
-
-private class KotlinPluginLifecycleCoroutineContextElement(
-    val lifecycle: KotlinPluginLifecycle,
-) : CoroutineContext.Element {
-    companion object Key : CoroutineContext.Key<KotlinPluginLifecycleCoroutineContextElement>
-
-    override val key: CoroutineContext.Key<KotlinPluginLifecycleCoroutineContextElement> = Key
-}
-
-private class RestrictedLifecycleStages(
-    private val lifecycle: KotlinPluginLifecycle,
-    private val allowedStages: Set<Stage>,
-) : CoroutineContext.Element, ContinuationInterceptor {
-
-    override val key: CoroutineContext.Key<*> = ContinuationInterceptor
-
-    override fun <T> interceptContinuation(continuation: Continuation<T>): Continuation<T> = object : Continuation<T> {
-        override val context: CoroutineContext
-            get() = continuation.context
-
-        override fun resumeWith(result: Result<T>) = when {
-            result.isFailure -> continuation.resumeWith(result)
-            lifecycle.stage !in allowedStages -> continuation.resumeWithException(
-                IllegalLifecycleException(
-                    "Required stage in '$allowedStages', but lifecycle switched to '${lifecycle.stage}'"
-                )
-            )
-            else -> continuation.resumeWith(result)
-        }
-    }
-
-    init {
-        if (lifecycle.stage !in allowedStages) {
-            throw IllegalLifecycleException("Required stage in '${allowedStages}' but lifecycle is currently in '${lifecycle.stage}'")
-        }
-    }
-}