[Gradle, BT] Add support for outputs backup within build tools API compilation

#KT-61896 Fixed

libraries/tools/kotlin-gradle-plugin-integration-tests/src/test/kotlin/org/jetbrains/kotlin/gradle/BuildToolsApiJvmCompilationIT.kt

@@ -13,6 +13,8 @@ import org.jetbrains.kotlin.gradle.tasks.KotlinCompilerExecutionStrategy
 import org.jetbrains.kotlin.gradle.tasks.USING_JVM_INCREMENTAL_COMPILATION_MESSAGE
 import org.jetbrains.kotlin.gradle.testbase.*
 import org.junit.jupiter.api.DisplayName
+import kotlin.io.path.deleteExisting
+import kotlin.io.path.writeText
 
 @DisplayName("JVM compilation via the Build Tools API")
 @JvmGradlePluginTests
@@ -55,6 +57,35 @@ class BuildToolsApiJvmCompilationIT : KGPBaseTest() {
         }
     }
 
+    @GradleTest
+    @DisplayName("Gradle side outputs backup works")
+    fun outputsBackupWorks(gradleVersion: GradleVersion) {
+        project(
+            "simpleProject", gradleVersion, buildOptions = defaultBuildOptions.copy(
+                usePreciseOutputsBackup = false,
+                keepIncrementalCompilationCachesInMemory = false,
+            )
+        ) {
+            build("compileKotlin") {
+                assertTasksExecuted(":compileKotlin")
+            }
+            val newBrokenSrc = kotlinSourcesDir().resolve("broken.kt")
+            newBrokenSrc.writeText(
+                //language=kt
+                """
+                broken code
+                """.trimIndent()
+            )
+            buildAndFail("compileKotlin") {
+                assertTasksFailed(":compileKotlin")
+            }
+            newBrokenSrc.deleteExisting()
+            build("compileKotlin") {
+                assertTasksUpToDate(":compileKotlin")
+            }
+        }
+    }
+
     @GradleTest
     @DisplayName("Simple project non-incremental in-process compilation")
     fun compileJvmInProcessNonIncremental(gradleVersion: GradleVersion) = testSimpleProject(

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/compilerRunner/GradleCompilerRunnerWithWorkers.kt

@@ -15,6 +15,7 @@ import org.gradle.workers.WorkParameters
 import org.gradle.workers.WorkQueue
 import org.gradle.workers.WorkerExecutor
 import org.jetbrains.kotlin.build.report.metrics.*
+import org.jetbrains.kotlin.gradle.logging.GradleKotlinLogger
 import org.jetbrains.kotlin.gradle.tasks.*
 import java.io.File
 import javax.inject.Inject
@@ -27,11 +28,11 @@ internal class GradleCompilerRunnerWithWorkers(
     jdkToolsJar: File?,
     compilerExecutionSettings: CompilerExecutionSettings,
     buildMetrics: BuildMetricsReporter<GradleBuildTime, GradleBuildPerformanceMetric>,
-    private val workerExecutor: WorkerExecutor
+    private val workerExecutor: WorkerExecutor,
 ) : GradleCompilerRunner(taskProvider, jdkToolsJar, compilerExecutionSettings, buildMetrics) {
     override fun runCompilerAsync(
         workArgs: GradleKotlinCompilerWorkArguments,
-        taskOutputsBackup: TaskOutputsBackup?
+        taskOutputsBackup: TaskOutputsBackup?,
     ): WorkQueue {
 
         buildMetrics.addTimeMetric(GradleBuildPerformanceMetric.CALL_WORKER)
@@ -49,10 +50,10 @@ internal class GradleCompilerRunnerWithWorkers(
     }
 
     internal abstract class GradleKotlinCompilerWorkAction @Inject constructor(
-        private val fileSystemOperations: FileSystemOperations
+        private val fileSystemOperations: FileSystemOperations,
     ) : WorkAction<GradleKotlinCompilerWorkParameters> {
 
-        private val logger = Logging.getLogger("kotlin-compile-worker")
+        private val logger = GradleKotlinLogger(Logging.getLogger("kotlin-compile-worker"))
 
         override fun execute() {
             val taskOutputsBackup = if (parameters.snapshotsDir.isPresent) {
@@ -61,7 +62,7 @@ internal class GradleCompilerRunnerWithWorkers(
                     parameters.buildDir,
                     parameters.snapshotsDir,
                     parameters.taskOutputsToRestore.get(),
-                    logger
+                    logger,
                 )
             } else {
                 null

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/compilerRunner/btapi/BuildToolsApiCompilationWork.kt

@@ -6,6 +6,7 @@
 package org.jetbrains.kotlin.compilerRunner.btapi
 
 import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.FileSystemOperations
 import org.gradle.api.provider.ListProperty
 import org.gradle.api.provider.Property
 import org.gradle.workers.WorkAction
@@ -25,13 +26,21 @@ import org.jetbrains.kotlin.gradle.logging.SL4JKotlinLogger
 import org.jetbrains.kotlin.gradle.plugin.BuildFinishedListenerService
 import org.jetbrains.kotlin.gradle.plugin.internal.BuildIdService
 import org.jetbrains.kotlin.gradle.plugin.internal.state.TaskLoggers
-import org.jetbrains.kotlin.gradle.tasks.KotlinCompilerExecutionStrategy
-import org.jetbrains.kotlin.gradle.tasks.throwExceptionIfCompilationFailed
+import org.jetbrains.kotlin.gradle.tasks.*
+import org.jetbrains.kotlin.gradle.tasks.CompilationErrorException
+import org.jetbrains.kotlin.gradle.tasks.FailedCompilationException
+import org.jetbrains.kotlin.gradle.tasks.OOMErrorException
+import org.jetbrains.kotlin.gradle.tasks.TaskOutputsBackup
 import org.jetbrains.kotlin.incremental.ClasspathChanges
 import org.slf4j.LoggerFactory
 import java.io.File
+import java.rmi.RemoteException
+import javax.inject.Inject
 
-internal abstract class BuildToolsApiCompilationWork : WorkAction<BuildToolsApiCompilationWork.BuildToolsApiCompilationParameters> {
+internal abstract class BuildToolsApiCompilationWork @Inject constructor(
+    private val fileSystemOperations: FileSystemOperations,
+) :
+    WorkAction<BuildToolsApiCompilationWork.BuildToolsApiCompilationParameters> {
     internal interface BuildToolsApiCompilationParameters : WorkParameters {
         val buildIdService: Property<BuildIdService>
         val buildFinishedListenerService: Property<BuildFinishedListenerService>
@@ -63,7 +72,7 @@ internal abstract class BuildToolsApiCompilationWork : WorkAction<BuildToolsApiC
             }
     }
 
-    override fun execute() {
+    private fun performCompilation(): CompilationResult {
         val executionStrategy = workArguments.compilerExecutionSettings.strategy
         try {
             val classLoader = parameters.classLoadersCachingService.get()
@@ -75,7 +84,9 @@ internal abstract class BuildToolsApiCompilationWork : WorkAction<BuildToolsApiC
             }
             val executionConfig = compilationService.makeCompilerExecutionStrategyConfiguration().apply {
                 when (executionStrategy) {
-                    KotlinCompilerExecutionStrategy.DAEMON -> useDaemonStrategy(workArguments.compilerExecutionSettings.daemonJvmArgs ?: emptyList())
+                    KotlinCompilerExecutionStrategy.DAEMON -> useDaemonStrategy(
+                        workArguments.compilerExecutionSettings.daemonJvmArgs ?: emptyList()
+                    )
                     KotlinCompilerExecutionStrategy.IN_PROCESS -> useInProcessStrategy()
                     else -> error("The \"$executionStrategy\" execution strategy is not supported by the Build Tools API")
                 }
@@ -110,19 +121,70 @@ internal abstract class BuildToolsApiCompilationWork : WorkAction<BuildToolsApiC
                     classpathSnapshotsConfig,
                 )
             }
-            val result = compilationService.compileJvm(
+            return compilationService.compileJvm(
                 buildId,
                 executionConfig,
                 jvmCompilationConfig,
                 emptyList(),
                 workArguments.compilerArgs.toList(),
             )
-            throwExceptionIfCompilationFailed(result.asExitCode, executionStrategy)
+        } catch (e: Throwable) {
+            if (e is OutOfMemoryError || e.hasOOMCause()) {
+                val helpMessage = when (executionStrategy) {
+                    KotlinCompilerExecutionStrategy.DAEMON -> kotlinDaemonOOMHelperMessage
+                    KotlinCompilerExecutionStrategy.IN_PROCESS -> kotlinInProcessOOMHelperMessage
+                    else -> error("The \"$executionStrategy\" execution strategy is not supported by the Build Tools API")
+                }
+                throw OOMErrorException(helpMessage)
+            } else if (e is RemoteException) {
+                throw DaemonCrashedException(e)
+            } else {
+                throw e
+            }
         } finally {
             log.info(executionStrategy.asFinishLogMessage)
         }
     }
 
+    // the files are backed up in the task action before any changes to the outputs
+    private fun initializeBackup(): TaskOutputsBackup? = if (parameters.snapshotsDir.isPresent) {
+        TaskOutputsBackup(
+            fileSystemOperations,
+            parameters.buildDir,
+            parameters.snapshotsDir,
+            parameters.taskOutputsToRestore.get(),
+            log,
+        )
+    } else {
+        null
+    }
+
+    override fun execute() {
+        val backup = initializeBackup()
+        val executionStrategy = workArguments.compilerExecutionSettings.strategy
+        try {
+            val result = performCompilation()
+            if (result == CompilationResult.COMPILATION_OOM_ERROR || result == CompilationResult.COMPILATION_ERROR) {
+                backup?.restoreOutputs()
+            }
+            throwExceptionIfCompilationFailed(result.asExitCode, executionStrategy)
+        } catch (e: FailedCompilationException) {
+            // Restore outputs only in cases where we expect that the user will make some changes to their project:
+            //   - For a compilation error, the user will need to fix their source code
+            //   - For an OOM error, the user will need to increase their memory settings
+            // In the other cases where there is nothing the user can fix in their project, we should not restore the outputs.
+            // Otherwise, the next build(s) will likely fail in exactly the same way as this build because their inputs and outputs are
+            // the same.
+            if (backup != null && (e is CompilationErrorException || e is OOMErrorException)) {
+                log.info("Restoring task outputs to pre-compilation state")
+                backup.restoreOutputs()
+            }
+            throw e
+        } finally {
+            backup?.deleteSnapshot()
+        }
+    }
+
     // temporary adapter property
     private val CompilationResult.asExitCode
         get() = when (this) {

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/tasks/AbstractKotlinCompile.kt

@@ -269,7 +269,7 @@ abstract class AbstractKotlinCompile<T : CommonCompilerArguments> @Inject constr
                             projectLayout.buildDirectory,
                             projectLayout.buildDirectory.dir("snapshot/kotlin/$name"),
                             outputsToRestore = allOutputFiles() - taskOutputsBackupExcludes.get(),
-                            logger
+                            GradleKotlinLogger(logger),
                         ).also {
                             it.createSnapshot()
                         }

libraries/tools/kotlin-gradle-plugin/src/common/kotlin/org/jetbrains/kotlin/gradle/tasks/TasksOutputsBackup.kt

@@ -8,8 +8,8 @@ package org.jetbrains.kotlin.gradle.tasks
 import org.gradle.api.file.Directory
 import org.gradle.api.file.DirectoryProperty
 import org.gradle.api.file.FileSystemOperations
-import org.gradle.api.logging.Logger
 import org.gradle.api.provider.Provider
+import org.jetbrains.kotlin.buildtools.api.KotlinLogger
 import java.io.File
 import java.net.URI
 import java.nio.file.FileSystems
@@ -34,7 +34,7 @@ internal class TaskOutputsBackup(
      */
     val outputsToRestore: List<File>,
 
-    val logger: Logger
+    val logger: KotlinLogger,
 ) {
 
     fun createSnapshot() {