[LL FIR] implement lock resolution algorithm for jumping phases

This commit extends `LLFirLockProvider` API functionality to provide `withJumpingLock`, which can be safely used in multithreaded scenarios. Comparing to regular `withWriteLock`/`withReadLock`, the new API has a bit more overhead as we have to maintain `jumpingResolutionStatesStack` and `FirInProcessOfResolvingToJumpingPhaseState` even without contention. See KDoc of `LLFirLockProvider#withJumpingLockImpl` for implementation details. The implicit type phase has been migrated to the new approach, so now it is safe to drop the last global phase lock. Potentially, we can have some benefits from migration of other jumping phases to this API. ^KT-56551 Fixed ^KT-55750 Fixed
2024-01-23 14:42:09 +01:00
parent 58ed52ee1f
commit e03f8b503f
5 changed files with 326 additions and 27 deletions
@@ -12,7 +12,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecific
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.lockWithPCECheck
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.*
-import java.util.concurrent.CountDownLatch
+import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater
 import java.util.concurrent.locks.ReentrantLock
@@ -75,6 +75,7 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
     * [action] will be executed once if [target] is not yet resolved to [phase] phase.
     *
     * @see withReadLock
+     * @see withJumpingLock
     */
    inline fun withWriteLock(
        target: FirElementWithResolveState,
@@ -104,16 +105,6 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
        }
    }

-    inline fun withJumpingLock(
-        target: FirElementWithResolveState,
-        phase: FirResolvePhase,
-        action: () -> Unit,
-    ) {
-        checker.lazyResolveToPhaseInside(phase, isJumpingPhase = true) {
-            target.withLock(toPhase = phase, updatePhase = true, action = action)
-        }
-    }
-
    /**
     * Locks an a [FirElementWithResolveState] to resolve from `toPhase - 1` to [toPhase] and
     * then updates the [FirElementWithResolveState.resolveState] to a [toPhase] if [updatePhase] is **true**.
@@ -178,6 +169,10 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh

                    return
                }
+
+                is FirInProcessOfResolvingToJumpingPhaseState -> {
+                    errorWithFirSpecificEntries("$stateSnapshot state are not allowed to be inside non-jumping lock", fir = this)
+                }
            }
        }
    }
@@ -192,8 +187,7 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
        toPhase: FirResolvePhase,
        stateSnapshot: FirResolveState,
    ) {
-        val latch = CountDownLatch(1)
-        val newState = FirInProcessOfResolvingToPhaseStateWithBarrier(toPhase, latch)
+        val newState = FirInProcessOfResolvingToPhaseStateWithBarrier(toPhase)
        resolveStateFieldUpdater.compareAndSet(this, stateSnapshot, newState)
    }

@@ -211,11 +205,182 @@ internal class LLFirLockProvider(private val checker: LLFirLazyResolveContractCh
            is FirInProcessOfResolvingToPhaseStateWithBarrier -> {
                stateSnapshotAfter.barrier.countDown()
            }
-            is FirResolvedToPhaseState -> {
+            is FirResolvedToPhaseState, is FirInProcessOfResolvingToJumpingPhaseState -> {
                errorWithFirSpecificEntries("phase is unexpectedly unlocked $stateSnapshotAfter", fir = this)
            }
        }
    }
+
+    /**
+     * Locks on an a [FirElementWithResolveState] to resolve from `phase - 1` to [phase] and
+     * then updates the [resolve state][FirElementWithResolveState.resolveState] to a [phase].
+     * Does nothing if [target] already has at least [phase] phase.
+     *
+     * @param actionUnderLock will be executed once under the lock if [target] is not yet resolved to [phase] phase and there are no cycles
+     * @param actionOnCycle will be executed once without the lock if [target] is not yet resolved to [phase] phase and a resolution cycle is found
+     *
+     * @see withWriteLock
+     * @see withJumpingLockImpl
+     */
+    fun withJumpingLock(
+        target: FirElementWithResolveState,
+        phase: FirResolvePhase,
+        actionUnderLock: () -> Unit,
+        actionOnCycle: () -> Unit,
+    ) {
+        checker.lazyResolveToPhaseInside(phase, isJumpingPhase = true) {
+            target.withJumpingLockImpl(phase, actionUnderLock, actionOnCycle)
+        }
+    }
+
+    /**
+     * Holds resolution states of the current thread.
+     * This information is required to properly process possible cycles
+     * during resolution.
+     *
+     * @see withJumpingLockImpl
+     * @see tryJumpingLock
+     * @see jumpingUnlock
+     */
+    private val jumpingResolutionStatesStack = JumpingResolutionStatesStack()
+
+    /**
+     * Locks an a [FirElementWithResolveState] to resolve from `toPhase - 1` to [toPhase] and
+     * then updates the [FirElementWithResolveState.resolveState] to a
+     * [toPhase] if no exceptions were found during [actionUnderLock].
+     *
+     * If [FirElementWithResolveState] is already at least at [toPhase], does nothing.
+     *
+     * ### Happy path:
+     *  1. Marks [FirElementWithResolveState] as in a process of resolve
+     *  2. Performs the resolve by calling [actionUnderLock]
+     *  3. Updates the resolve phase to [toPhase] if there is no exceptions
+     *  4. Notifies other threads waiting on the same lock that this thread already resolved the declaration,
+     *  so other threads can continue its execution
+     *
+     *  ### Cycle handling
+     *  During step 1 we can realize someone already set [FirInProcessOfResolvingToJumpingPhaseState]
+     *  for the current [FirElementWithResolveState], so there is a room for a possible deadlock.
+     *
+     *  The requirement for the deadlock is not empty [jumpingResolutionStatesStack] as we should already hold another lock.
+     *  Otherwise, we can just wait on the [latch][FirInProcessOfResolvingToJumpingPhaseState.latch].
+     *
+     *  In the case of not empty [jumpingResolutionStatesStack], we have the following algorithm:
+     *  1. Set [waitingFor][FirInProcessOfResolvingToJumpingPhaseState.waitingFor] for the previous state
+     *  as we have an intention to take the next lock
+     *  2. Iterate over all [waitingFor][FirInProcessOfResolvingToJumpingPhaseState.waitingFor] recursively
+     *  to detect the possible cycle
+     *  3. Execute [actionOnCycle] without the lock in the case of cycle or waining on
+     *  the [latch][FirInProcessOfResolvingToJumpingPhaseState.latch] to try to take the lock again later
+     *
+     * @param actionUnderLock will be executed once under the lock if [this] is not yet resolved to [toPhase] phase and there are no cycles
+     * @param actionOnCycle will be executed once without the lock if [this] is not yet resolved to [toPhase] phase and a resolution cycle is found
+     *
+     *  @see withJumpingLock
+     */
+    private fun FirElementWithResolveState.withJumpingLockImpl(
+        toPhase: FirResolvePhase,
+        actionUnderLock: () -> Unit,
+        actionOnCycle: () -> Unit,
+    ) {
+        while (true) {
+            checkCanceled()
+
+            @OptIn(ResolveStateAccess::class)
+            val currentState = resolveState
+            if (currentState.resolvePhase >= toPhase) {
+                // already resolved by some other thread
+                return
+            }
+
+            when (currentState) {
+                is FirResolvedToPhaseState -> {
+                    if (!tryJumpingLock(toPhase, currentState)) continue
+
+                    var exceptionOccurred = false
+                    try {
+                        actionUnderLock()
+                    } catch (e: Throwable) {
+                        exceptionOccurred = true
+                        throw e
+                    } finally {
+                        val newPhase = if (!exceptionOccurred) toPhase else currentState.resolvePhase
+                        jumpingUnlock(toPhase = newPhase)
+                    }
+
+                    return
+                }
+
+                is FirInProcessOfResolvingToJumpingPhaseState -> {
+                    val previousState = jumpingResolutionStatesStack.peek()
+
+                    // Not null value means we already hold a lock for another declaration in the current thread,
+                    // so we have to check the possible cycle
+                    if (previousState != null) {
+                        // All writing to waitingFor will be consistent, as it is the last writing if we have cycle
+                        previousState.waitingFor = currentState
+
+                        // Cycle check
+                        var nextState: FirInProcessOfResolvingToJumpingPhaseState? = currentState
+                        while (nextState != null) {
+                            if (nextState === previousState) {
+                                previousState.waitingFor = null
+                                return actionOnCycle()
+                            }
+
+                            nextState = nextState.waitingFor
+                        }
+                    }
+
+                    try {
+                        // Waiting until another thread released the lock
+                        currentState.latch.await(DEFAULT_LOCKING_INTERVAL, TimeUnit.MILLISECONDS)
+                    } finally {
+                        previousState?.waitingFor = null
+                    }
+                }
+
+                is FirInProcessOfResolvingToPhaseStateWithoutBarrier, is FirInProcessOfResolvingToPhaseStateWithBarrier -> {
+                    errorWithFirSpecificEntries("$currentState state are not allowed to be inside jumping lock", fir = this)
+                }
+            }
+        }
+    }
+
+    /**
+     * Trying to set [FirInProcessOfResolvingToJumpingPhaseState] to [this].
+     *
+     * @return **true** if the state is published successfully
+     *
+     * @see withJumpingLockImpl
+     * @see FirInProcessOfResolvingToJumpingPhaseState
+     */
+    private fun FirElementWithResolveState.tryJumpingLock(
+        toPhase: FirResolvePhase,
+        stateSnapshot: FirResolveState,
+    ): Boolean {
+        val newState = FirInProcessOfResolvingToJumpingPhaseState(toPhase)
+        val isSucceed = resolveStateFieldUpdater.compareAndSet(this, stateSnapshot, newState)
+        if (!isSucceed) return false
+
+        jumpingResolutionStatesStack.push(newState)
+
+        return true
+    }
+
+    /**
+     * Publish [FirResolvedToPhaseState] with [toPhase] phase and unlocks current [FirInProcessOfResolvingToJumpingPhaseState].
+     *
+     * @see withJumpingLockImpl
+     * @see FirInProcessOfResolvingToJumpingPhaseState
+     * @see FirResolvedToPhaseState
+     */
+    private fun FirElementWithResolveState.jumpingUnlock(toPhase: FirResolvePhase) {
+        val currentState = jumpingResolutionStatesStack.pop()
+
+        resolveStateFieldUpdater.set(this, FirResolvedToPhaseState(toPhase))
+        currentState.latch.countDown()
+    }
 }

 private val resolveStateFieldUpdater = AtomicReferenceFieldUpdater.newUpdater(
@@ -229,7 +394,61 @@ private val globalLockEnabled: Boolean by lazy(LazyThreadSafetyMode.PUBLICATION)
 }

 private val implicitPhaseLockEnabled: Boolean by lazy(LazyThreadSafetyMode.PUBLICATION) {
-    Registry.`is`("kotlin.implicit.resolve.phase.under.global.lock", true)
+    Registry.`is`("kotlin.implicit.resolve.phase.under.global.lock", false)
 }

-private const val DEFAULT_LOCKING_INTERVAL = 50L
+private const val DEFAULT_LOCKING_INTERVAL = 50L
+
+/**
+ * @see FirInProcessOfResolvingToJumpingPhaseState
+ */
+private class JumpingResolutionStatesStack {
+    private val stateStackHolder = ThreadLocal.withInitial<MutableList<FirInProcessOfResolvingToJumpingPhaseState>> {
+        mutableListOf()
+    }
+
+    /**
+     * Adds [newState] to the stack and set [waitingFor][FirInProcessOfResolvingToJumpingPhaseState.waitingFor]
+     * for the previous state if needed
+     */
+    fun push(newState: FirInProcessOfResolvingToJumpingPhaseState) {
+        val states = stateStackHolder.get()
+
+        val currentState = states.lastOrNull()
+        currentState?.waitingFor = newState
+        states += newState
+    }
+
+    /**
+     * Pops from the top of the stack the last state and return it.
+     * Updates [waitingFor][FirInProcessOfResolvingToJumpingPhaseState.waitingFor] for
+     * the previous state if needed
+     *
+     * Note: it doesn't release the [lock][FirInProcessOfResolvingToJumpingPhaseState.latch]
+     */
+    fun pop(): FirInProcessOfResolvingToJumpingPhaseState {
+        val states = stateStackHolder.get()
+
+        val currentState = states.removeLast()
+        val prevState = states.lastOrNull()
+        requireWithAttachment(
+            condition = prevState == null || prevState.waitingFor === currentState,
+            message = { "The lock contact is violated" },
+        )
+
+        prevState?.waitingFor = null
+
+        // Drop the empty stack to avoid memory leak
+        // as the updated capacity of the stack can be high
+        if (states.isEmpty()) {
+            stateStackHolder.remove()
+        }
+
+        return currentState
+    }
+
+    /**
+     * Current state on the top if exists
+     */
+    fun peek(): FirInProcessOfResolvingToJumpingPhaseState? = stateStackHolder.get().lastOrNull()
+}
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -26,9 +26,11 @@ import org.jetbrains.kotlin.fir.symbols.FirBasedSymbol
 import org.jetbrains.kotlin.fir.symbols.impl.FirCallableSymbol
 import org.jetbrains.kotlin.fir.types.FirImplicitTypeRef
 import org.jetbrains.kotlin.fir.util.setMultimapOf
+import org.jetbrains.kotlin.fir.utils.exceptions.withFirEntry
 import org.jetbrains.kotlin.fir.utils.exceptions.withFirSymbolEntry
 import org.jetbrains.kotlin.fir.visitors.transformSingle
 import org.jetbrains.kotlin.utils.exceptions.errorWithAttachment
+import org.jetbrains.kotlin.utils.exceptions.requireWithAttachment

 internal object LLFirImplicitTypesLazyResolver : LLFirLazyResolver(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
    override fun resolve(
@@ -105,6 +107,29 @@ internal class LLImplicitBodyResolveComputationSession : ImplicitBodyResolveComp
    fun postponedSymbols(target: FirCallableDeclaration): Collection<FirBasedSymbol<*>> {
        return postponedSymbols[target.symbol]
    }
+
+    private var cycledSymbol: FirCallableSymbol<*>? = null
+
+    /**
+     * Push [symbol] with a recursion return type to be able to report it later
+     *
+     * @param symbol is a symbol with the recursion error in the return type
+     *
+     * @see popCycledSymbolIfExists
+     * @see LLFirImplicitBodyTargetResolver.handleCycleInResolution
+     */
+    fun pushCycledSymbol(symbol: FirCallableSymbol<*>) {
+        requireWithAttachment(cycledSymbol == null, { "Nested recursion is not allowed" })
+        cycledSymbol = symbol
+    }
+
+    /**
+     * Pop [FirCallableSymbol] with a recursion return type if it was [pushed][pushCycledSymbol]
+     *
+     * @see pushCycledSymbol
+     * @see org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirReturnTypeCalculatorWithJump.resolveDeclaration
+     */
+    fun popCycledSymbolIfExists(): FirCallableSymbol<*>? = cycledSymbol?.also { cycledSymbol = null }
 }

 internal class LLFirImplicitBodyTargetResolver(
@@ -139,6 +164,17 @@ internal class LLFirImplicitBodyTargetResolver(
        }
    }

+    /**
+     * @see org.jetbrains.kotlin.analysis.low.level.api.fir.element.builder.LLFirReturnTypeCalculatorWithJump.resolveDeclaration
+     */
+    override fun handleCycleInResolution(target: FirElementWithResolveState) {
+        requireWithAttachment(target is FirCallableDeclaration, { "Resolution cycle is supposed to be only for callable declaration" }) {
+            withFirEntry("target", target)
+        }
+
+        llImplicitBodyResolveComputationSession.pushCycledSymbol(target.symbol)
+    }
+
    override fun doLazyResolveUnderLock(target: FirElementWithResolveState) {
        when {
            target is FirCallableDeclaration && target.isCopyCreatedInScope -> {
@@ -1,5 +1,5 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

@@ -35,18 +35,24 @@ internal class LLFirReturnTypeCalculatorWithJump(

        val designation = declaration.collectDesignationWithFile().asResolveTarget()
        val targetSession = designation.target.llFirSession
+        val computationSession = implicitBodyResolveComputationSession as LLImplicitBodyResolveComputationSession
        val resolver = LLFirImplicitBodyTargetResolver(
            designation,
            lockProvider = lockProvider,
            scopeSession = targetSession.getScopeSession(),
            firResolveContextCollector = towerDataContextCollector,
-            llImplicitBodyResolveComputationSessionParameter = implicitBodyResolveComputationSession as LLImplicitBodyResolveComputationSession,
+            llImplicitBodyResolveComputationSessionParameter = computationSession,
        )

        lockProvider.withGlobalPhaseLock(FirResolvePhase.IMPLICIT_TYPES_BODY_RESOLVE) {
            resolver.resolveDesignation()
        }

+        // Report recursion error if we found cycle during resolution
+        if (computationSession.popCycledSymbolIfExists() == declaration.symbol) {
+            return recursionInImplicitTypeRef()
+        }
+
        LLFirImplicitTypesLazyResolver.checkIsResolved(designation)
        return declaration.returnTypeRef as FirResolvedTypeRef
    }
@@ -9,6 +9,7 @@ import org.jetbrains.kotlin.analysis.low.level.api.fir.api.targets.*
 import org.jetbrains.kotlin.analysis.low.level.api.fir.file.builder.LLFirLockProvider
 import org.jetbrains.kotlin.analysis.low.level.api.fir.sessions.LLFirSession
 import org.jetbrains.kotlin.analysis.low.level.api.fir.util.checkPhase
+import org.jetbrains.kotlin.analysis.low.level.api.fir.util.errorWithFirSpecificEntries
 import org.jetbrains.kotlin.fir.FirElement
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.FirFileAnnotationsContainer
@@ -172,9 +173,12 @@ internal abstract class LLFirTargetResolver(
            lockProvider.withJumpingLock(
                target,
                resolverPhase,
-                action = {
+                actionUnderLock = {
                    doLazyResolveUnderLock(target)
                    updatePhaseForDeclarationInternals(target)
+                },
+                actionOnCycle = {
+                    handleCycleInResolution(target)
                }
            )
        } else {
@@ -184,6 +188,19 @@ internal abstract class LLFirTargetResolver(
        }
    }

+    /**
+     * Will be executed in the case of detected cycle between elements during jumping resolve.
+     *
+     * **There is no guaranties that [target] is guarded by the lock of the current thread**
+     *
+     * @param target an element with detected cycle
+     *
+     * @see LLFirLockProvider.withJumpingLock
+     */
+    protected open fun handleCycleInResolution(target: FirElementWithResolveState) {
+        errorWithFirSpecificEntries("Resolution cycle is detected", fir = target)
+    }
+
    /**
     * Execute [action] under the write lock in the context of [target].
     *
@@ -1,14 +1,14 @@
 /*
- * Copyright 2010-2023 JetBrains s.r.o. and Kotlin Programming Language contributors.
+ * Copyright 2010-2024 JetBrains s.r.o. and Kotlin Programming Language contributors.
 * Use of this source code is governed by the Apache 2.0 license that can be found in the license/LICENSE.txt file.
 */

 package org.jetbrains.kotlin.fir.declarations

-import java.util.concurrent.CountDownLatch
 import org.jetbrains.kotlin.fir.FirElementWithResolveState
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticProperty
 import org.jetbrains.kotlin.fir.declarations.synthetic.FirSyntheticPropertyAccessor
+import java.util.concurrent.CountDownLatch

 /**
 * The current lazy resolve state of some [org.jetbrains.kotlin.fir.FirElementWithResolveState].
@@ -99,10 +99,9 @@ class FirInProcessOfResolvingToPhaseStateWithoutBarrier private constructor(
 *
 * @see FirResolveState
 */
-class FirInProcessOfResolvingToPhaseStateWithBarrier(
-    override val resolvingTo: FirResolvePhase,
-    val barrier: CountDownLatch,
-) : FirInProcessOfResolvingToPhaseState() {
+class FirInProcessOfResolvingToPhaseStateWithBarrier(override val resolvingTo: FirResolvePhase) : FirInProcessOfResolvingToPhaseState() {
+    val barrier: CountDownLatch = CountDownLatch(1)
+
    init {
        require(resolvingTo != FirResolvePhase.RAW_FIR) {
            "Cannot resolve to ${FirResolvePhase.RAW_FIR} as it's a first phase"
@@ -110,4 +109,26 @@ class FirInProcessOfResolvingToPhaseStateWithBarrier(
    }

    override fun toString(): String = "ResolvingToWithBarrier($resolvingTo)"
-}
+}
+
+/**
+ * The class representing the lazy resolve state of some [FirElementWithResolveState] in a case
+ * when some thread is resolving it from [resolvePhase] to [resolvingTo] and potentially can have a cycle
+ * between threads.
+ *
+ * Some other threads can wait on a [latch].
+ *
+ * [waitingFor] shows that the current state is waining for the result of another [FirElementWithResolveState] element.
+ * This another element can be in the process of resolution as from the same thread and also from another thread.
+ * In the second case the current thread will wait on the corresponding [latch].
+ *
+ * @see FirResolveState
+ */
+class FirInProcessOfResolvingToJumpingPhaseState(override val resolvingTo: FirResolvePhase) : FirInProcessOfResolvingToPhaseState() {
+    val latch = CountDownLatch(1)
+
+    @Volatile
+    var waitingFor: FirInProcessOfResolvingToJumpingPhaseState? = null
+
+    override fun toString(): String = "ResolvingJumpingTo($resolvingTo)"
+}