fix(build-service,tooling-api): Gradle Daemon and ToolingAPI server is not stopped properly

The Tooling API server (which is executed in a separate process) was not being stopped properly after IToolingApiServer.shutdown() was called. As a result, the Gradle daemon started by the tooling API was not stopped as well.

This commit fixes this issue by properly implementing the shutdown() call in Tooling API. Now, when the user closes a project and the GradleBuildService is destroyed, the Tooling API server stops the Gradle daemon(s) and then shutdowns itself.

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

@@ -31,6 +31,7 @@ import java.nio.charset.StandardCharsets;
 import java.util.HashSet;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
 import org.gradle.tooling.ConfigurableLauncher;
@@ -59,8 +60,26 @@ public class Main {
     LOG.debug("Server started. Will run until shutdown message is received...");
     try {
       Main.future.get();
+    } catch (CancellationException cancellationException) {
+      // ignored
     } catch (InterruptedException | ExecutionException e) {
       LOG.error("An error occurred while waiting for shutdown message", e);
+    } finally {
+
+      // Cleanup should be performed in ToolingApiServerImpl.shutdown()
+      // this is to make sure that the daemons are stopped in case the client doesn't call shutdown()
+      try {
+        if (server.isInitialized() || server.isConnected()) {
+          LOG.warn("Connection to tooling server closed without shutting it down!");
+          server.shutdown().get();
+        }
+      } catch (InterruptedException | ExecutionException e) {
+        e.printStackTrace();
+      } finally {
+        future = null;
+        client = null;
+        JvmLogger.interceptor = null;
+      }
     }
   }
 

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

@@ -50,12 +50,17 @@ import org.gradle.tooling.CancellationTokenSource
 import org.gradle.tooling.GradleConnectionException
 import org.gradle.tooling.GradleConnector
 import org.gradle.tooling.ProjectConnection
+import org.gradle.tooling.ResultHandler
 import org.gradle.tooling.UnsupportedVersionException
 import org.gradle.tooling.exceptions.UnsupportedBuildArgumentException
 import org.gradle.tooling.exceptions.UnsupportedOperationConfigurationException
+import org.gradle.tooling.internal.consumer.BlockingResultHandler
+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.atomic.AtomicBoolean
+import kotlin.system.exitProcess
 
 /**
  * Implementation for the Gradle Tooling API server.
@@ -65,7 +70,7 @@ import java.util.concurrent.CompletionException
 internal class ToolingApiServerImpl(private val project: ProjectImpl) :
   IToolingApiServer {
 
-  private var initialized = false
+  private val _isInitialized = AtomicBoolean(false)
   private var client: IToolingApiClient? = null
   private var connector: GradleConnector? = null
   private var connection: ProjectConnection? = null
@@ -73,6 +78,31 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
   private var buildCancellationToken: CancellationTokenSource? = null
   private val log = ILogger.newInstance(javaClass.simpleName)
 
+  /**
+   * Whether the project has been initialized or not.
+   */
+  var isInitialized: Boolean
+    get() = _isInitialized.get()
+    private set(value) = _isInitialized.set(value)
+
+  /**
+   * Whether the server has a live connection to Gradle.
+   */
+  val isConnected: Boolean
+    get() = connector != null || connection != null
+
+  companion object {
+
+    /**
+     * Time duration for which the the Tooling API server waits after calling
+     * [DefaultGradleConnector.close] and before exiting the server's process.
+     *
+     * This delay should be long enough to let the tooling API stop the daemon but short enough so
+     * that the server's process is not kept alive for longer duration.
+     */
+    const val DELAY_BEFORE_EXIT_MS = 1000L
+  }
+
   override fun initialize(params: InitializeProjectParams): CompletableFuture<InitializeResult> {
     return CompletableFuture.supplyAsync {
       try {
@@ -133,7 +163,7 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
 
         this.project.rootProject = project.rootProject
         this.project.projects = project.projects
-        initialized = true
+        this.isInitialized = true
 
         notifyBuildSuccess(emptyList())
         return@supplyAsync InitializeResult(true)
@@ -146,7 +176,7 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
   }
 
   override fun isServerInitialized(): CompletableFuture<Boolean> {
-    return CompletableFuture.supplyAsync { initialized }
+    return CompletableFuture.supplyAsync { isInitialized }
   }
 
   override fun getRootProject(): CompletableFuture<IProject> {
@@ -257,21 +287,38 @@ internal class ToolingApiServerImpl(private val project: ProjectImpl) :
   }
 
   override fun shutdown(): CompletableFuture<Void> {
-    return CompletableFuture.runAsync {
-      connection?.close()
-      connector?.disconnect()
-      Main.future?.cancel(true)
+    log.info("Shutting down Tooling API Server...")
 
-      this.connection = null
-      this.connector = null
-      this.client = null
-      this.buildCancellationToken = null
-      this.lastInitParams = null
-      Main.future = null
-      Main.client = null
+    connection?.close()
+    connector?.disconnect()
+    connection = null
+    connector = null
 
-      this.initialized = false
+    // Stop all daemons
+    log.info("Stopping all Gradle Daemons...")
+    DefaultGradleConnector.close()
+
+    try {
+      // wait for the daemon to be stopped properly
+      Thread.sleep(DELAY_BEFORE_EXIT_MS)
+    } catch (e: Exception) {
+      // ignored
     }
+
+    // update the initialization flag before cancelling future
+    this.isInitialized = false
+
+    Main.future?.cancel(true)
+
+    this.client = null
+    this.buildCancellationToken = null // connector.disconnect() cancels any running builds
+    this.lastInitParams = null
+    Main.future = null
+    Main.client = null
+
+    // no need to return anything
+    // tooling server should be restarted once it has been shutdown
+    exitProcess(0)
   }
 
   private fun getTaskFailureType(error: Throwable): Failure =

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

@@ -22,7 +22,6 @@ import com.itsaky.androidide.tooling.api.messages.TaskExecutionMessage
 import com.itsaky.androidide.tooling.api.messages.result.BuildCancellationRequestResult
 import com.itsaky.androidide.tooling.api.messages.result.InitializeResult
 import com.itsaky.androidide.tooling.api.messages.result.TaskExecutionResult
-import org.eclipse.lsp4j.jsonrpc.services.JsonNotification
 import org.eclipse.lsp4j.jsonrpc.services.JsonRequest
 import org.eclipse.lsp4j.jsonrpc.services.JsonSegment
 import java.util.concurrent.*
@@ -36,13 +35,16 @@ import java.util.concurrent.*
 interface IToolingApiServer {
 
   /** Initialize the server with the project directory. */
-  @JsonRequest fun initialize(params: InitializeProjectParams): CompletableFuture<InitializeResult>
+  @JsonRequest
+  fun initialize(params: InitializeProjectParams): CompletableFuture<InitializeResult>
 
   /** Is the server initialized? */
-  @JsonRequest fun isServerInitialized(): CompletableFuture<Boolean>
+  @JsonRequest
+  fun isServerInitialized(): CompletableFuture<Boolean>
 
   /** Get the root project. */
-  @JsonRequest fun getRootProject(): CompletableFuture<IProject>
+  @JsonRequest
+  fun getRootProject(): CompletableFuture<IProject>
 
   /** Execute the tasks specified in the message. */
   @JsonRequest
@@ -54,12 +56,14 @@ interface IToolingApiServer {
    * @return A [CompletableFuture] which completes when the current build cancellation process
    * finishes (either successfully or with an error).
    */
-  @JsonRequest fun cancelCurrentBuild(): CompletableFuture<BuildCancellationRequestResult>
+  @JsonRequest
+  fun cancelCurrentBuild(): CompletableFuture<BuildCancellationRequestResult>
 
   /**
-   * Shutdown the tooling API server. This will disconnect all the project connection instances.]
+   * Shutdown the tooling API server. This will disconnect all the project connection instances.
    *
    * @return A [CompletableFuture] which completes when the shutdown process is finished.
    */
-  @JsonNotification fun shutdown(): CompletableFuture<Void>
+  @JsonRequest
+  fun shutdown(): CompletableFuture<Void>
 }