IDEA-321510 [Gradle|Sync] fix: Gradle sync can be canceled during handling Gradle call completion.

GitOrigin-RevId: f12b77ddd6e7febe461bb557c8088d589fd4bbb4

plugins/gradle/src/org/jetbrains/plugins/gradle/service/buildActionRunner/GradleBuildActionResultHandler.kt

@@ -65,18 +65,17 @@ class GradleBuildActionResultHandler(
 
   fun createProjectLoadedHandler(): IntermediateResultHandler<GradleModelHolderState> {
     return IntermediateResultHandler { state ->
-      try {
+      runCancellable {
         onProjectLoaded(state)
       }
-      catch (e: ProcessCanceledException) {
-        resolverCtx.cancel()
-      }
     }
   }
 
   fun createBuildFinishedHandler(): IntermediateResultHandler<GradleModelHolderState> {
     return IntermediateResultHandler { state ->
-      onBuildCompleted(state)
+      runCancellable {
+        onBuildCompleted(state)
+      }
     }
   }
 
@@ -91,7 +90,9 @@ class GradleBuildActionResultHandler(
       override fun onComplete(result: Any?) {
         try {
           if (result is GradleModelHolderState) {
-            onBuildCompleted(result)
+            runCancellable {
+              onBuildCompleted(result)
+            }
           }
         }
         finally {
@@ -101,7 +102,9 @@ class GradleBuildActionResultHandler(
 
       override fun onFailure(failure: GradleConnectionException) {
         try {
-          onBuildFailed(failure)
+          runCancellable {
+            onBuildFailed(failure)
+          }
         }
         finally {
           buildFinishWaiter.countDown()
@@ -109,4 +112,15 @@ class GradleBuildActionResultHandler(
       }
     }
   }
+
+  private fun runCancellable(action: () -> Unit) {
+    try {
+      if (!resolverCtx.isCancellationRequested) {
+        action()
+      }
+    }
+    catch (e: ProcessCanceledException) {
+      resolverCtx.cancel()
+    }
+  }
 }

plugins/gradle/testSources/org/jetbrains/plugins/gradle/service/buildActionRunner/GradlePhasedSyncTest.kt

@@ -1,6 +1,7 @@
 // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
 package org.jetbrains.plugins.gradle.service.buildActionRunner
 
+import com.intellij.openapi.progress.ProcessCanceledException
 import com.intellij.openapi.util.Disposer
 import com.intellij.openapi.util.use
 import org.jetbrains.plugins.gradle.importing.BuildFinishedModel
@@ -9,6 +10,8 @@ import org.jetbrains.plugins.gradle.importing.ProjectLoadedModel
 import org.jetbrains.plugins.gradle.importing.ProjectLoadedModelProvider
 import org.jetbrains.plugins.gradle.tooling.builder.ProjectPropertiesTestModelBuilder
 import org.junit.Test
+import org.junit.jupiter.api.Assertions
+import java.util.concurrent.atomic.AtomicBoolean
 
 class GradlePhasedSyncTest : GradlePhasedSyncTestCase() {
 
@@ -71,4 +74,52 @@ class GradlePhasedSyncTest : GradlePhasedSyncTestCase() {
       buildCompletionAssertion.assertListenerState(1) { "Build action should be finished only once" }
     }
   }
+
+  @Test
+  fun `test one-phased Gradle sync cancellation`() {
+    Disposer.newDisposable().use { disposable ->
+      val isSyncCancelled = AtomicBoolean(false)
+
+      whenBuildCompleted(disposable) {
+        throw ProcessCanceledException()
+      }
+
+      createSettingsFile("")
+      importProject(errorHandler = { _, _ ->
+        isSyncCancelled.set(true)
+      })
+
+      Assertions.assertTrue(isSyncCancelled.get()) {
+        "Gradle sync should be cancelled."
+      }
+    }
+  }
+
+  @Test
+  fun `test two-phased Gradle sync cancellation`() {
+    Disposer.newDisposable().use { disposable ->
+      val isBuildFinished = AtomicBoolean(false)
+      val isSyncCancelled = AtomicBoolean(false)
+
+      whenProjectLoaded(disposable) {
+        throw ProcessCanceledException()
+      }
+      whenBuildCompleted(disposable) {
+        isBuildFinished.set(true)
+      }
+
+      createSettingsFile("")
+      importProject(errorHandler = { _, _ ->
+        isSyncCancelled.set(true)
+      })
+
+      Assertions.assertTrue(isSyncCancelled.get()) {
+        "Gradle sync should be cancelled during the project loaded action."
+      }
+      Assertions.assertFalse(isBuildFinished.get()) {
+        "Gradle sync should be cancelled during the project loaded action.\n" +
+        "Therefore the project build action shouldn't be completed."
+      }
+    }
+  }
 }

plugins/gradle/testSources/org/jetbrains/plugins/gradle/service/buildActionRunner/GradlePhasedSyncTestCase.kt

@@ -15,13 +15,36 @@ import org.jetbrains.plugins.gradle.service.project.GradleProjectResolverExtensi
 import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
 import java.util.concurrent.CopyOnWriteArrayList
 import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.AtomicReference
 
 abstract class GradlePhasedSyncTestCase : GradleImportingTestCase() {
 
+  private lateinit var syncErrorHandler: AtomicReference<(String, String?) -> Unit>
+  private val defaultSyncErrorHandler = { errorMessage: String, errorDetails: String? ->
+    super.handleImportFailure(errorMessage, errorDetails)
+  }
+
   override fun setUp() {
     super.setUp()
+
     myProject.registerServiceInstance(TestProjectResolverExtensionService::class.java, TestProjectResolverExtensionService())
     GradleProjectResolverExtension.EP_NAME.point.registerExtension(TestProjectResolverExtension(), testRootDisposable)
+
+    syncErrorHandler = AtomicReference(defaultSyncErrorHandler)
+  }
+
+  fun importProject(errorHandler: (String, String?) -> Unit) {
+    syncErrorHandler.set(errorHandler)
+    try {
+      importProject()
+    }
+    finally {
+      syncErrorHandler.set(defaultSyncErrorHandler)
+    }
+  }
+
+  override fun handleImportFailure(errorMessage: String, errorDetails: String?) {
+    syncErrorHandler.get()(errorMessage, errorDetails)
   }
 
   fun addToolingExtensionClasses(parentDisposable: Disposable, vararg toolingExtensionClasses: Class<*>) {