[K/N] In executors make sure child process does not leak.

native/executors/src/main/kotlin/org/jetbrains/kotlin/native/executors/HostExecutor.kt

@@ -7,15 +7,127 @@
 
 package org.jetbrains.kotlin.native.executors
 
-import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.*
-import kotlinx.coroutines.Job
-import kotlinx.coroutines.launch
-import kotlinx.coroutines.runBlocking
 import java.io.File
+import java.io.IOException
+import java.io.InputStream
+import java.io.OutputStream
+import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
 import java.util.logging.Logger
 import kotlin.time.*
 
+class ProcessStreams(
+    process: Process,
+    stdin: InputStream,
+    stdout: OutputStream,
+    stderr: OutputStream,
+    jobLauncher: (suspend () -> Unit) -> Job,
+) {
+    private val ignoreIOErrors = AtomicBoolean(false)
+    private val stdin = jobLauncher {
+        stdin.apply {
+            copyStreams(this, process.outputStream)
+            close()
+        }
+        process.outputStream.close()
+    }
+    private val stdout = jobLauncher {
+        stdout.apply {
+            copyStreams(process.inputStream, this)
+            close()
+        }
+        process.inputStream.close()
+    }
+    private val stderr = jobLauncher {
+        stderr.apply {
+            copyStreams(process.errorStream, this)
+            close()
+        }
+        process.errorStream.close()
+    }
+
+    private fun copyStreams(from: InputStream, to: OutputStream) {
+        try {
+            from.copyTo(to)
+        } catch(e: IOException) {
+            if (ignoreIOErrors.get())
+                return
+            throw e
+        }
+    }
+
+    suspend fun drain() {
+        // First finish passing input into the process.
+        stdin.join()
+        // Now receive all the output in whatever order.
+        stdout.join()
+        stderr.join()
+    }
+
+    fun cancel() {
+        ignoreIOErrors.set(true)
+        stdout.cancel()
+        stderr.cancel()
+        stdin.cancel()
+    }
+}
+
+fun CoroutineScope.pumpStreams(
+    process: Process,
+    stdin: InputStream,
+    stdout: OutputStream,
+    stderr: OutputStream,
+) = ProcessStreams(
+    process,
+    stdin,
+    stdout,
+    stderr,
+) {
+    launch {
+        it()
+    }
+}
+
+private object ProcessKiller {
+    init {
+        Runtime.getRuntime().addShutdownHook(Thread {
+            killAllProcesses()
+        })
+    }
+
+    private var processes = ConcurrentLinkedQueue<Process>()
+
+    private fun killAllProcesses() {
+        processes.forEach {
+            it.destroyForcibly()
+        }
+    }
+
+    fun register(process: Process) = processes.add(process)
+
+    fun deregister(process: Process) = processes.remove(process)
+}
+
+private fun <T> ProcessBuilder.scoped(block: suspend CoroutineScope.(Process) -> T): T {
+    val process = start()
+    // Make sure the process is killed even if the jvm process is being destroyed.
+    // e.g. gradle --no-daemon task execution was cancelled by the user pressing ^C
+    ProcessKiller.register(process)
+    return try {
+        runBlocking(Dispatchers.IO) {
+            block(process)
+        }
+    } finally {
+        // Make sure the process is killed even if the current thread was interrupted.
+        // e.g. gradle task execution was cancelled by the user pressing ^C
+        process.destroyForcibly()
+        // The process is dead, no need to ensure its destruction during the shutdown.
+        ProcessKiller.deregister(process)
+    }
+}
+
 /**
  * [Executor] that runs the process on the host system.
  */
@@ -23,60 +135,37 @@ class HostExecutor : Executor {
     private val logger = Logger.getLogger(HostExecutor::class.java.name)
 
     override fun execute(request: ExecuteRequest): ExecuteResponse {
-        return runBlocking(Dispatchers.IO) {
+        val workingDirectory = request.workingDirectory ?: File(request.executableAbsolutePath).parentFile
-            val workingDirectory = request.workingDirectory ?: File(request.executableAbsolutePath).parentFile
+        val commandLine = "${request.executableAbsolutePath}${request.args.joinToString(separator = " ", prefix = " ")}"
-            val commandLine = "${request.executableAbsolutePath}${request.args.joinToString(separator = " ", prefix = " ")}"
+        val environmentFormatted =
-            logger.info("""
+            request.environment.entries.joinToString(prefix = "{", postfix = "}") { "\"${it.key}\": \"${it.value}\"" }
+        logger.info(
+            """
                 |Starting command: $commandLine
                 |In working directory: ${workingDirectory.absolutePath}
-                |With additional environment: ${request.environment.entries.joinToString(prefix = "{", postfix = "}") { "\"${it.key}\": \"${it.value}\"" }}
+                |With additional environment: $environmentFormatted
                 |And timeout: ${request.timeout}
-                """.trimMargin())
+                """.trimMargin()
-            val exitCodeWithTime = measureTimedValue {
+        )
-                val process = ProcessBuilder(listOf(request.executableAbsolutePath) + request.args).apply {
+        return ProcessBuilder(listOf(request.executableAbsolutePath) + request.args).apply {
-                    directory(workingDirectory)
+            directory(workingDirectory)
-                    environment().putAll(request.environment)
+            environment().putAll(request.environment)
-                }.start()
+        }.scoped { process ->
-                val jobs: MutableList<Job> = mutableListOf()
+            val streams = pumpStreams(process, request.stdin, request.stdout, request.stderr)
-                jobs.add(launch {
+            val (isTimeout, duration) = measureTimedValue {
-                    request.stdin.apply {
+                !process.waitFor(request.timeout.inWholeMilliseconds, TimeUnit.MILLISECONDS)
-                        copyTo(process.outputStream)
+            }
-                        close()
+            if (isTimeout) {
-                    }
+                logger.warning("Timeout running $commandLine in $duration")
-                    process.outputStream.close()
+                streams.cancel()
-                })
+                process.destroyForcibly()
-                jobs.add(launch {
+                streams.drain()
-                    request.stdout.apply {
+                ExecuteResponse(null, duration)
-                        process.inputStream.copyTo(this)
+            } else {
-                        close()
+                logger.info("Finished executing $commandLine in $duration exit code ${process.exitValue()}")
-                    }
+                streams.drain()
-                    process.inputStream.close()
+                ExecuteResponse(process.exitValue(), duration)
-                })
-                jobs.add(launch {
-                    request.stderr.apply {
-                        process.errorStream.copyTo(this)
-                        close()
-                    }
-                    process.errorStream.close()
-                })
-
-                if (!process.waitFor(request.timeout.inWholeMilliseconds, TimeUnit.MILLISECONDS)) {
-                    logger.warning("Timeout running $commandLine")
-                    // Cancel every stream, no need to wait for them.
-                    jobs.forEach {
-                        it.cancel()
-                    }
-                    process.destroyForcibly()
-                    null
-                } else {
-                    // Drain every stream.
-                    jobs.forEach {
-                        it.join()
-                    }
-                    process.exitValue()
-                }
             }
-            ExecuteResponse(exitCodeWithTime.value, exitCodeWithTime.duration)
         }
     }
 }