Perform tryLock-and-checkCanceled on waiting lock in LockBasedStorageManager

Relates to #KT-38012

compiler/frontend/src/org/jetbrains/kotlin/context/context.kt

@@ -16,6 +16,7 @@
 
 package org.jetbrains.kotlin.context
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.builtins.KotlinBuiltIns
 import org.jetbrains.kotlin.descriptors.ModuleDescriptor
@@ -85,7 +86,9 @@ class MutableModuleContextImpl(
 
 fun GlobalContext(debugName: String): GlobalContextImpl {
     val tracker = ExceptionTracker()
-    return GlobalContextImpl(LockBasedStorageManager.createWithExceptionHandling(debugName, tracker), tracker)
+    return GlobalContextImpl(LockBasedStorageManager.createWithExceptionHandling(debugName, tracker) {
+        ProgressManager.checkCanceled()
+    }, tracker)
 }
 
 fun ProjectContext(project: Project, debugName: String): ProjectContext = ProjectContextImpl(project, GlobalContext(debugName))

core/util.runtime/src/org/jetbrains/kotlin/storage/LockBasedStorageManager.java

@@ -29,8 +29,6 @@ import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
-import java.util.concurrent.locks.Lock;
-import java.util.concurrent.locks.ReentrantLock;
 
 public class LockBasedStorageManager implements StorageManager {
     private static final String PACKAGE_NAME = StringsKt.substringBeforeLast(LockBasedStorageManager.class.getCanonicalName(), ".", "");
@@ -55,7 +53,7 @@ public class LockBasedStorageManager implements StorageManager {
         RuntimeException handleException(@NotNull Throwable throwable);
     }
 
-    public static final StorageManager NO_LOCKS = new LockBasedStorageManager("NO_LOCKS", ExceptionHandlingStrategy.THROW, NoLock.INSTANCE) {
+    public static final StorageManager NO_LOCKS = new LockBasedStorageManager("NO_LOCKS", ExceptionHandlingStrategy.THROW, EmptySimpleLock.INSTANCE) {
         @NotNull
         @Override
         protected <T> RecursionDetectedResult<T> recursionDetectedDefault() {
@@ -64,18 +62,31 @@ public class LockBasedStorageManager implements StorageManager {
     };
 
     @NotNull
-    public static LockBasedStorageManager createWithExceptionHandling(@NotNull String debugText, @NotNull ExceptionHandlingStrategy exceptionHandlingStrategy) {
+    public static LockBasedStorageManager createWithExceptionHandling(
-        return new LockBasedStorageManager(debugText, exceptionHandlingStrategy, new ReentrantLock());
+            @NotNull String debugText,
+            @NotNull ExceptionHandlingStrategy exceptionHandlingStrategy
+    ) {
+        return createWithExceptionHandling(debugText, exceptionHandlingStrategy, null);
     }
 
-    protected final Lock lock;
+    @NotNull
+    public static LockBasedStorageManager createWithExceptionHandling(
+            @NotNull String debugText,
+            @NotNull ExceptionHandlingStrategy exceptionHandlingStrategy,
+            @Nullable Runnable checkCancelled
+    ) {
+        return new LockBasedStorageManager(debugText, exceptionHandlingStrategy,
+                                           SimpleLock.Companion.simpleLock(checkCancelled));
+    }
+
+    protected final SimpleLock lock;
     private final ExceptionHandlingStrategy exceptionHandlingStrategy;
     private final String debugText;
 
     private LockBasedStorageManager(
             @NotNull String debugText,
             @NotNull ExceptionHandlingStrategy exceptionHandlingStrategy,
-            @NotNull Lock lock
+            @NotNull SimpleLock lock
     ) {
         this.lock = lock;
         this.exceptionHandlingStrategy = exceptionHandlingStrategy;
@@ -83,7 +94,11 @@ public class LockBasedStorageManager implements StorageManager {
     }
 
     public LockBasedStorageManager(String debugText) {
-        this(debugText, ExceptionHandlingStrategy.THROW, new ReentrantLock());
+        this(debugText, null);
+    }
+
+    public LockBasedStorageManager(String debugText, @Nullable Runnable checkCancelled) {
+        this(debugText, ExceptionHandlingStrategy.THROW, SimpleLock.Companion.simpleLock(checkCancelled));
     }
 
     @Override

core/util.runtime/src/org/jetbrains/kotlin/storage/locks.kt

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2010-2020 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.storage
+
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.locks.Lock
+import java.util.concurrent.locks.ReentrantLock
+
+private const val CHECK_CANCELLATION_PERIOD_MS: Long = 50
+
+interface SimpleLock {
+    fun lock()
+
+    fun unlock()
+
+    companion object {
+        fun simpleLock(checkCancelled: Runnable? = null) =
+            checkCancelled?.let { CancellableSimpleLock(it) } ?: DefaultSimpleLock()
+    }
+}
+
+inline fun <T> SimpleLock.guarded(crossinline computable: () -> T): T {
+    lock()
+    return try {
+        computable()
+    } finally {
+        unlock()
+    }
+}
+
+object EmptySimpleLock : SimpleLock {
+    override fun lock() {
+    }
+
+    override fun unlock() {
+    }
+}
+
+open class DefaultSimpleLock(protected val lock: Lock = ReentrantLock()) : SimpleLock {
+
+    override fun lock() = lock.lock()
+
+    override fun unlock() = lock.unlock()
+
+}
+
+class CancellableSimpleLock(lock: Lock, private val checkCancelled: Runnable) : DefaultSimpleLock(lock) {
+    constructor(checkCancelled: Runnable) : this(checkCancelled = checkCancelled, lock = ReentrantLock())
+
+    override fun lock() {
+        while (!lock.tryLock(CHECK_CANCELLATION_PERIOD_MS, TimeUnit.MILLISECONDS)) {
+            //ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+            checkCancelled.run()
+        }
+    }
+
+}

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/PerFileAnalysisCache.kt

@@ -38,6 +38,8 @@ import org.jetbrains.kotlin.resolve.diagnostics.Diagnostics
 import org.jetbrains.kotlin.resolve.diagnostics.DiagnosticsElementsCache
 import org.jetbrains.kotlin.resolve.lazy.BodyResolveMode
 import org.jetbrains.kotlin.resolve.lazy.ResolveSession
+import org.jetbrains.kotlin.storage.CancellableSimpleLock
+import org.jetbrains.kotlin.storage.guarded
 import org.jetbrains.kotlin.types.KotlinType
 import org.jetbrains.kotlin.util.slicedMap.ReadOnlySlice
 import org.jetbrains.kotlin.util.slicedMap.WritableSlice
@@ -54,6 +56,9 @@ internal class PerFileAnalysisCache(val file: KtFile, componentProvider: Compone
     private val cache = HashMap<PsiElement, AnalysisResult>()
     private var fileResult: AnalysisResult? = null
     private val lock = ReentrantLock()
+    private val guardLock = CancellableSimpleLock(lock) {
+        ProgressIndicatorProvider.checkCanceled()
+    }
 
     internal fun fetchAnalysisResults(element: KtElement): AnalysisResult? {
         assert(element.containingKtFile == file) { "Wrong file. Expected $file, but was ${element.containingKtFile}" }
@@ -75,28 +80,23 @@ internal class PerFileAnalysisCache(val file: KtFile, componentProvider: Compone
 
         val analyzableParent = KotlinResolveDataProvider.findAnalyzableParent(element)
 
-        lock.lock()
+        return guardLock.guarded {
-        try {
-            ProgressIndicatorProvider.checkCanceled()
-
             // step 1: perform incremental analysis IF it is applicable
-            getIncrementalAnalysisResult()?.let { return it }
+            getIncrementalAnalysisResult()?.let { return@guarded it }
 
             // cache does not contain AnalysisResult per each kt/psi element
             // instead it looks up analysis for its parents - see lookUp(analyzableElement)
 
             // step 2: return result if it is cached
             lookUp(analyzableParent)?.let {
-                return it
+                return@guarded it
             }
 
             // step 3: perform analyze of analyzableParent as nothing has been cached yet
             val result = analyze(analyzableParent)
             cache[analyzableParent] = result
 
-            return result
+            return@guarded result
-        } finally {
-            lock.unlock()
         }
     }
 

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/ProjectResolutionFacade.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.idea.caches.resolve
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiElement
 import com.intellij.psi.util.CachedValueProvider
@@ -20,6 +21,8 @@ import org.jetbrains.kotlin.idea.caches.trackers.KotlinCodeBlockModificationList
 import org.jetbrains.kotlin.psi.KtElement
 import org.jetbrains.kotlin.psi.KtFile
 import org.jetbrains.kotlin.resolve.CompositeBindingContext
+import org.jetbrains.kotlin.storage.CancellableSimpleLock
+import org.jetbrains.kotlin.storage.guarded
 import org.jetbrains.kotlin.utils.addToStdlib.firstNotNullResult
 import java.util.concurrent.locks.ReentrantLock
 
@@ -48,6 +51,10 @@ internal class ProjectResolutionFacade(
         get() = globalContext.storageManager.compute { cachedValue.value }
 
     private val analysisResultsLock = ReentrantLock()
+    private val analysisResultsSimpleLock = CancellableSimpleLock(analysisResultsLock) {
+        ProgressManager.checkCanceled()
+    }
+
     private val analysisResults = CachedValuesManager.getManager(project).createCachedValue(
         {
             val resolverForProject = cachedResolverForProject
@@ -142,13 +149,8 @@ internal class ProjectResolutionFacade(
     internal fun getAnalysisResultsForElements(elements: Collection<KtElement>): AnalysisResult {
         assert(elements.isNotEmpty()) { "elements collection should not be empty" }
 
-        val slruCache = run {
+        val slruCache = analysisResultsSimpleLock.guarded {
-            analysisResultsLock.lock()
+            analysisResults.value!!
-            try {
-                analysisResults.value!!
-            } finally {
-                analysisResultsLock.unlock()
-            }
         }
         val results = elements.map {
             val perFileCache = slruCache[it.containingKtFile]

idea/idea-analysis/src/org/jetbrains/kotlin/idea/caches/resolve/util/globalContextUtils.kt

@@ -9,6 +9,7 @@ import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.context.GlobalContextImpl
 import org.jetbrains.kotlin.idea.caches.resolve.PlatformAnalysisSettings
 import org.jetbrains.kotlin.idea.project.useCompositeAnalysis
+import org.jetbrains.kotlin.progress.ProgressIndicatorAndCompilationCanceledStatus
 import org.jetbrains.kotlin.storage.ExceptionTracker
 import org.jetbrains.kotlin.storage.LockBasedStorageManager
 
@@ -23,7 +24,9 @@ private fun GlobalContextImpl.contextWithCompositeExceptionTracker(debugName: St
 private fun GlobalContextImpl.contextWithNewLockAndCompositeExceptionTracker(debugName: String): GlobalContextImpl {
     val newExceptionTracker = CompositeExceptionTracker(this.exceptionTracker)
     return GlobalContextImpl(
-        LockBasedStorageManager.createWithExceptionHandling(debugName, newExceptionTracker),
+        LockBasedStorageManager.createWithExceptionHandling(debugName, newExceptionTracker) {
+            ProgressIndicatorAndCompilationCanceledStatus.checkCanceled()
+        },
         newExceptionTracker
     )
 }

plugins/scripting/scripting-compiler-impl/src/org/jetbrains/kotlin/scripting/resolve/scriptAnnotationsPreprocessing.kt

@@ -5,6 +5,7 @@
 
 package org.jetbrains.kotlin.scripting.resolve
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import org.jetbrains.kotlin.builtins.DefaultBuiltIns
 import org.jetbrains.kotlin.config.LanguageVersionSettingsImpl
@@ -27,7 +28,13 @@ internal fun String?.orAnonymous(kind: String = ""): String =
         this ?: "<anonymous" + (if (kind.isNotBlank()) " $kind" else "") + ">"
 
 internal fun constructAnnotation(psi: KtAnnotationEntry, targetClass: KClass<out Annotation>, project: Project): Annotation {
-    val module = ModuleDescriptorImpl(Name.special("<script-annotations-preprocessing>"), LockBasedStorageManager("scriptAnnotationsPreprocessing"), DefaultBuiltIns.Instance)
+    val module = ModuleDescriptorImpl(
+        Name.special("<script-annotations-preprocessing>"),
+        LockBasedStorageManager("scriptAnnotationsPreprocessing") {
+            ProgressManager.checkCanceled()
+        },
+        DefaultBuiltIns.Instance
+    )
     val evaluator = ConstantExpressionEvaluator(module, LanguageVersionSettingsImpl.DEFAULT, project)
     val trace = BindingTraceContext()