fix(tooling-api): add cancellation support for project initialization

subprojects/tooling-api-impl/src/main/java/com/itsaky/androidide/tooling/impl/ToolingApiServerImpl.kt

@@ -46,6 +46,7 @@ import com.itsaky.androidide.tooling.api.messages.result.TaskExecutionResult.Fai
 import com.itsaky.androidide.tooling.impl.internal.ProjectImpl
 import com.itsaky.androidide.tooling.impl.sync.ModelBuilderException
 import com.itsaky.androidide.tooling.impl.sync.RootModelBuilder
+import com.itsaky.androidide.tooling.impl.sync.RootProjectModelBuilderParams
 import com.itsaky.androidide.tooling.impl.util.StopWatch
 import com.itsaky.androidide.utils.ILogger
 import org.gradle.tooling.BuildCancelledException
@@ -61,6 +62,8 @@ import org.gradle.tooling.internal.consumer.DefaultGradleConnector
 import java.io.File
 import java.util.concurrent.CompletableFuture
 import java.util.concurrent.CompletionException
+import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
 import kotlin.system.exitProcess
 
 /**
@@ -75,9 +78,14 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
   private var connector: GradleConnector? = null
   private var connection: ProjectConnection? = null
   private var lastInitParams: InitializeProjectParams? = null
-  private var buildCancellationToken: CancellationTokenSource? = null
+  private var _buildCancellationToken: CancellationTokenSource? = null
   private val log = ILogger.newInstance(javaClass.simpleName)
 
+  private val cancellationTokenAccessLock = ReentrantLock(/* fair = */ true)
+  private var buildCancellationToken: CancellationTokenSource?
+    get() = cancellationTokenAccessLock.withLock { _buildCancellationToken }
+    set(value) = cancellationTokenAccessLock.withLock { _buildCancellationToken = value }
+
   /**
    * Whether the project has been initialized or not.
    */
@@ -162,8 +170,14 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
 
         stopWatch.lapFromLast("Project connection established")
 
+        this.buildCancellationToken = GradleConnector.newCancellationTokenSource()
+
         val project = try {
-          val impl = RootModelBuilder(params).build(connection) as? ProjectImpl?
+          val modelBuilderParams = RootProjectModelBuilderParams(
+            connection,
+            this.buildCancellationToken!!.token()
+          )
+          val impl = RootModelBuilder(params).build(modelBuilderParams) as? ProjectImpl?
             ?: throw ModelBuilderException("Failed to build project model")
           impl
         } catch (err: Throwable) {
@@ -309,6 +323,7 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
 
       try {
         this.buildCancellationToken!!.cancel()
+        this.buildCancellationToken = null
       } catch (e: Exception) {
         val failureReason = CANCELLATION_ERROR
         failureReason.message = "${failureReason.message}: ${e.message}"

subprojects/tooling-api-impl/src/main/java/com/itsaky/androidide/tooling/impl/sync/RootModelBuilder.kt

@@ -29,7 +29,6 @@ import com.itsaky.androidide.tooling.impl.Main.finalizeLauncher
 import com.itsaky.androidide.tooling.impl.internal.ProjectImpl
 import com.itsaky.androidide.utils.ILogger
 import org.gradle.tooling.ConfigurableLauncher
-import org.gradle.tooling.ProjectConnection
 import org.gradle.tooling.model.idea.IdeaProject
 import java.io.Serializable
 
@@ -39,16 +38,18 @@ import java.io.Serializable
  * @author Akash Yadav
  */
 class RootModelBuilder(initializationParams: InitializeProjectParams) :
-  AbstractModelBuilder<ProjectConnection, IProject>(initializationParams), Serializable {
+  AbstractModelBuilder<RootProjectModelBuilderParams, IProject>(initializationParams), Serializable {
 
   private val serialVersionUID = 1L
 
-  override fun build(param: ProjectConnection): IProject {
+  override fun build(param: RootProjectModelBuilderParams): IProject {
+
+    val (projectConnection, cancellationToken) = param
 
     // do not reference the 'initializationParams' field in the
     val initializationParams = initializationParams
 
-    val executor = param.action { controller ->
+    val executor = projectConnection.action { controller ->
       val ideaProject = controller.getModelAndLog(IdeaProject::class.java)
 
       val ideaModules = ideaProject.modules
@@ -106,6 +107,10 @@ class RootModelBuilder(initializationParams: InitializeProjectParams) :
     finalizeLauncher(executor)
     applyAndroidModelBuilderProps(executor)
 
+    if (cancellationToken != null) {
+      executor.withCancellationToken(cancellationToken)
+    }
+
     val logger = ILogger.newInstance("ProjectReader")
     logger.warn("Starting build. See build output for more details...")
 

subprojects/tooling-api-impl/src/main/java/com/itsaky/androidide/tooling/impl/sync/types.kt

@@ -19,9 +19,22 @@ package com.itsaky.androidide.tooling.impl.sync
 
 import com.android.builder.model.v2.models.Versions
 import org.gradle.tooling.BuildController
+import org.gradle.tooling.CancellationToken
+import org.gradle.tooling.ProjectConnection
 import org.gradle.tooling.model.idea.IdeaModule
 import org.gradle.tooling.model.idea.IdeaProject
 
+/**
+ * Parameters for the root project model builder.
+ *
+ * @property projectConnection The project connection
+ * @property cancellationToken The cancellation token.
+ */
+data class RootProjectModelBuilderParams(
+  val projectConnection: ProjectConnection,
+  val cancellationToken: CancellationToken?
+)
+
 /**
  * Parameters for building model for an Android project.
  *

subprojects/tooling-api-impl/src/test/java/com/itsaky/androidide/tooling/impl/MultiModuleAndroidProjectTest.kt

@@ -141,6 +141,25 @@ class MultiModuleAndroidProjectTest {
     )
   }
 
+  @Test
+  fun `test project initialization cancellation`() {
+    // launch project initialization
+    val (server, _, initFuture) = ToolingApiTestLauncher().launchServerAsync()
+    Thread.sleep(1000L)
+
+    // cancel initialization request
+    val cancellationResult = server.cancelCurrentBuild().get()
+    println("Cancellation result: $cancellationResult")
+    assertThat(cancellationResult).isNotNull()
+    assertThat(cancellationResult.wasEnqueued).isTrue()
+    assertThat(cancellationResult.failureReason).isNull()
+
+    // verify that the initialization failed with reason BUILD_CANCELLED
+    val initResult = initFuture.get()
+    assertThat(initResult!!.isSuccessful).isFalse()
+    assertThat(initResult.failure).isEqualTo(TaskExecutionResult.Failure.BUILD_CANCELLED)
+  }
+
   private fun doAssertions(project: IProject, server: IToolingApiServer) {
     assertThat(project).isNotNull()
 

testing/tooling/src/main/java/com/itsaky/androidide/tooling/testing/ToolingApiTestLauncher.kt

@@ -66,7 +66,7 @@ class ToolingApiTestLauncher {
       "jdk.compiler" to "com.sun.tools.javac.util")
 
   @JvmOverloads
-  fun launchServer(
+  fun launchServerAsync(
     projectDir: Path = FileProvider.testProjectRoot(),
     client: MultiVersionTestClient = MultiVersionTestClient(),
     initParams: InitializeProjectParams = InitializeProjectParams(
@@ -76,7 +76,7 @@ class ToolingApiTestLauncher {
     log: ILogger = ILogger.newInstance("BuildOutputLogger"),
     sysProps: Map<String, String> = emptyMap(),
     sysEnvs: Map<String, String> = emptyMap()
-  ): Triple<IToolingApiServer, IProject, InitializeResult?> {
+  ): Triple<IToolingApiServer, IProject, CompletableFuture<InitializeResult>> {
     val cmdLine = createProcessCmd(
       FileProvider.implModule()
         .resolve("build/libs/tooling-api-all.jar").pathString,
@@ -101,10 +101,27 @@ class ToolingApiTestLauncher {
 
     val server = launcher.remoteProxy as IToolingApiServer
     val project = launcher.remoteProxy as IProject
-    val result = server.initialize(initParams).get()
+    val result = server.initialize(initParams)
     return Triple(server, project, result)
   }
 
+  @JvmOverloads
+  fun launchServer(
+    projectDir: Path = FileProvider.testProjectRoot(),
+    client: MultiVersionTestClient = MultiVersionTestClient(),
+    initParams: InitializeProjectParams = InitializeProjectParams(
+      projectDir.pathString,
+      client.gradleDistParams
+    ),
+    log: ILogger = ILogger.newInstance("BuildOutputLogger"),
+    sysProps: Map<String, String> = emptyMap(),
+    sysEnvs: Map<String, String> = emptyMap()
+  ): Triple<IToolingApiServer, IProject, InitializeResult?> {
+    return launchServerAsync(projectDir, client, initParams, log, sysProps, sysEnvs).let { result ->
+      Triple(result.first, result.second, result.third.get())
+    }
+  }
+
   private fun createProcessCmd(
     jar: String,
     sysProps: Map<String, String> = emptyMap()