From ca88e6e834aff3c79c3feae95272298239bdca38 Mon Sep 17 00:00:00 2001
From: Vladimir Dolzhenko <vladimir.dolzhenko@jetbrains.com>
Date: Mon, 4 Mar 2024 16:57:32 +0000
Subject: [PATCH] Try to acquire LC lock with a smaller time window with
 checkCancels

#KTIJ-28688


Merge-request: KT-MR-14735
Merged-by: Vladimir Dolzhenko <Vladimir.Dolzhenko@jetbrains.com>
---
 .../asJava/classes/LightClassesLazyCreator.kt | 31 +++++++++++++------
 1 file changed, 22 insertions(+), 9 deletions(-)

diff --git a/analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/classes/LightClassesLazyCreator.kt b/analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/classes/LightClassesLazyCreator.kt
index a0eacf74668..84feb417158 100644
--- a/analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/classes/LightClassesLazyCreator.kt
+++ b/analysis/light-classes-base/src/org/jetbrains/kotlin/asJava/classes/LightClassesLazyCreator.kt
@@ -5,12 +5,15 @@
 
 package org.jetbrains.kotlin.asJava.classes
 
+import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.project.Project
 import com.intellij.psi.PsiManager
 import com.intellij.psi.impl.PsiCachedValueImpl
 import com.intellij.psi.util.CachedValueProvider
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.locks.ReentrantLock
+import kotlin.time.DurationUnit
+import kotlin.time.toDuration
 
 class LightClassesLazyCreator(private val project: Project) : KotlinClassInnerStuffCache.LazyCreator() {
     override fun <T : Any> get(initializer: () -> T, dependencies: List<Any>) = object : Lazy<T> {
@@ -42,20 +45,30 @@ class LightClassesLazyCreator(private val project: Project) : KotlinClassInnerSt
 
                     // TODO: NOTE: acquire lock for a several seconds to avoid dead-lock via resolve is a WORKAROUND
 
-                    if (!initIsRunning.get() && lock.tryLock(5, TimeUnit.SECONDS)) {
-                        try {
-                            initIsRunning.set(true)
+                    val timeout = 5.toDuration(DurationUnit.SECONDS)
+                    val stepTimeout = 500.toDuration(DurationUnit.MILLISECONDS)
+
+                    var calculatedValue: T? = null
+                    for (attempt in 0 until (timeout.inWholeMilliseconds / stepTimeout.inWholeMilliseconds)) {
+                        if (calculatedValue != null || initIsRunning.get()) break
+
+                        ProgressManager.checkCanceled()
+
+                        if (lock.tryLock(stepTimeout.inWholeMilliseconds, TimeUnit.MILLISECONDS)) {
                             try {
-                                computeValue()
+                                initIsRunning.set(true)
+                                try {
+                                    calculatedValue = computeValue()
+                                } finally {
+                                    initIsRunning.set(false)
+                                }
                             } finally {
-                                initIsRunning.set(false)
+                                lock.unlock()
                             }
-                        } finally {
-                            lock.unlock()
                         }
-                    } else {
-                        computeValue()
                     }
+
+                    calculatedValue ?: computeValue()
                 }
             }