Wait on input stream redirection to finish

Avoids race where output is still copying after process exits
PaperMC · Oct 10, 2023 · 5389e61 · 5389e61
1 parent b1011c7
commit 5389e61
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 12 deletions.
diff --git a/paperweight-lib/src/main/kotlin/io/papermc/paperweight/tasks/GenerateDevBundle.kt b/paperweight-lib/src/main/kotlin/io/papermc/paperweight/tasks/GenerateDevBundle.kt
@@ -300,13 +300,17 @@ abstract class GenerateDevBundle : DefaultTask() {
             .directory(dir)
             .start()
 
-        val errBytes = ByteArrayOutputStream().also { redirect(process.errorStream, it) }
-        val outBytes = ByteArrayOutputStream().also { redirect(process.inputStream, it) }
+        val errBytes = ByteArrayOutputStream()
+        val errFuture = redirect(process.errorStream, errBytes)
+        val outBytes = ByteArrayOutputStream()
+        val outFuture = redirect(process.inputStream, outBytes)
 
         if (!process.waitFor(10L, TimeUnit.SECONDS)) {
             process.destroyForcibly()
             throw PaperweightException("Command '${cmd.joinToString(" ")}' did not finish after 10 seconds, killed process")
         }
+        errFuture.get(500L, TimeUnit.MILLISECONDS)
+        outFuture.get(500L, TimeUnit.MILLISECONDS)
         val err = asString(errBytes)
         val exit = process.exitValue()
         if (exit != 0 && exit != 1 || err.isNotBlank()) {

diff --git a/paperweight-lib/src/main/kotlin/io/papermc/paperweight/util/jar-runner.kt b/paperweight-lib/src/main/kotlin/io/papermc/paperweight/util/jar-runner.kt
@@ -24,6 +24,7 @@ package io.papermc.paperweight.util
 
 import io.papermc.paperweight.PaperweightException
 import java.io.OutputStream
+import java.util.concurrent.TimeUnit
 import java.util.jar.JarFile
 import kotlin.io.path.*
 import org.gradle.api.file.FileCollection
@@ -77,10 +78,12 @@ fun JavaLauncher.runJar(
     val process = processBuilder.start()
 
     output.use {
-        redirect(process.inputStream, it)
-        redirect(process.errorStream, it)
+        val outFuture = redirect(process.inputStream, it)
+        val errFuture = redirect(process.errorStream, it)
 
         val exit = process.waitFor()
+        outFuture.get(500L, TimeUnit.MILLISECONDS)
+        errFuture.get(500L, TimeUnit.MILLISECONDS)
         if (exit != 0) {
             val logMsg = logFilePath?.let { p -> " Log file: ${p.absolutePathString()}" } ?: ""
             throw PaperweightException("Execution of '$mainClass' failed with exit code $exit.$logMsg Classpath: ${classpath.asPath}")

diff --git a/paperweight-lib/src/main/kotlin/io/papermc/paperweight/util/utils.kt b/paperweight-lib/src/main/kotlin/io/papermc/paperweight/util/utils.kt
@@ -44,6 +44,7 @@ import java.util.Collections
 import java.util.IdentityHashMap
 import java.util.Locale
 import java.util.Optional
+import java.util.concurrent.CompletableFuture
 import java.util.concurrent.ThreadLocalRandom
 import java.util.concurrent.atomic.AtomicLong
 import kotlin.io.path.*
@@ -130,18 +131,20 @@ val Project.isBaseExecution: Boolean
 
 val redirectThreadCount: AtomicLong = AtomicLong(0)
 
-fun redirect(input: InputStream, out: OutputStream): Thread {
-    return Thread {
+fun redirect(input: InputStream, out: OutputStream): CompletableFuture<Unit> {
+    val future = CompletableFuture<Unit>()
+    val thread = Thread {
         try {
             input.copyTo(out)
-        } catch (e: Exception) {
-            throw PaperweightException("", e)
+            future.complete(Unit)
+        } catch (e: Throwable) {
+            future.completeExceptionally(PaperweightException("Failed to copy $input to $out", e))
         }
-    }.apply {
-        name = "paperweight stream redirect thread #${redirectThreadCount.getAndIncrement()}"
-        isDaemon = true
-        start()
     }
+    thread.name = "paperweight stream redirect thread #${redirectThreadCount.getAndIncrement()}"
+    thread.isDaemon = true
+    thread.start()
+    return future
 }
 
 object UselessOutputStream : OutputStream() {